AutoDIR: Automatic All-in-One Image Restoration with Latent Diffusion 论文阅读笔记
-
这是ECCV2024的论文,作者单位是是港中文和上海AI Lab
-
文章提出了一个叫AutoDIR的方法,包括两个关键阶段,一个是BIQA,基于vision-language model (CLIP范式)进行图像降质检测,一个是All-in-One 的 image restoration模型,使用 latent diffusion来处理多种图像降质。此外,本文提出的方法还支持用户定制。文章代码开源:https://github.com/jiangyitong/AutoDIR
-
整体流程如下图所示:
-
先预设一些句子,通过替换句子中的单词,每个句子其实针对了一种降质;然后这些句子经过文本编码器,得到text embedding;图像经过图像编码器,得到image embedding,和这些text embedding计算相似度,然后softmax来加权平均这些text embedding,就得到对这张图片的降质的描述的text embedding。用这个混合的embedding送进第二阶段的diffusion模型进行restoration,即可实现All-in-One;而用户定制则可以将这个text embedding替换成用户的文本对应的text embedding。
-
文章 提出,直接用分类器,即使是用ViT,来对图片进行降质分类,准确率也只有77.65%,很难做得好。而pretrain的CLIP会更注重语义而不注重图像质量。所以需要在质量标注上finetune一下。比如有5个针对不同降质的数据集,那么就有5个label和5个text embedding,再加上他们的GT是 no artifact,一共6个label 6个text embedding,可以和图像进行对比学习,计算图像提取的embedding和这6个文本embedding之间的余弦相似度,再过softmax,得到概率分布,计算多分类损失,fix住text encoder,就可以finetune image encoder,从而得到一个用于提取图像质量embedding的image encoder。
-
损失函数如上所示,这个y在对应的i下为1,否则为0,所以求和其实只有一项,就是对应降质的概率越高越好,也就是对应的embedding相似度越高越好。文章发现仅有这个损失还不够,所以加了下面这个损失:
-
这个损失的意思是,某种降质图片的gt图像,其对应降质的概率应该最低。
-
可以看到,通过这种方式训练的CLIP,其对降质的感知能力强,分类准确率能到97.94%:
-
restoration网络用的是LDM,但是LDM用的VAE,有压缩导致图像质量并不是特别好,更严重的问题是图像内容有畸变,跟原图不严格一致,这和restoration任务是不匹配的。为此文章提出了一个叫structure-correction module的模块,把降质图片和stable diffusion生成的图片concate起来送进这个模块,输出一个残差来修正stable diffusion生成结果的内容畸变。这里w是个可调的权重,LLIE任务里面就大一点,要求内容一致性高一点,SR任务就小一点。附录里面提到这个模块是个NAFNet。此外,这里的stable diffusion是要finetune的。
-
从实验结果上看,在单一任务上的指标也就一般,甚至PSNR垫底:
-
不过在All-in-One上的对比会领先: