AutoDIR: Automatic All-in-One Image Restoration with Latent Diffusion论文阅读笔记

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上的对比会领先：