AutoDIR: Automatic All-in-One Image Restoration with Latent Diffusion

AutoDIR: Automatic All-in-One Image Restoration with Latent Diffusion (Paper reading)

Yitong Jiang, The Chinese University of Hong Kong, arXiv23, Code, Paper

1. 前言

我们提出了一种具有潜在扩散的一体化图像恢复系统，名为AutoDIR，它可以自动检测和恢复具有多种未知退化的图像。我们的主要假设是，许多图像恢复任务，如超分辨率、运动去模糊、去噪、弱光增强、去雾和去噪，通常可以分解为一些共同的基算子，这些算子可以在不同方向上提高图像质量。AutoDIR旨在通过与多个图像恢复任务的联合训练，学习一个能够执行这些基本算子的统一图像恢复模型。具体而言，AutoDIR由基于CLIP的盲图像质量评估（BIQA）模块、基于潜在扩散的多功能图像恢复（AIR）模块和结构校正模块组成，BIQA模块自动检测输入图像的未知图像退化，AIR模块处理多种类型的图像退化，SCM模块进一步恢复图像结构。广泛的实验评估表明，AutoDIR在更广泛的图像恢复任务中优于最先进的方法。AutoDIR的设计还实现了灵活的用户控制（通过文本提示）和对新任务的泛化，作为图像恢复的基础模型。

2. 整体思想

整体思想就是分类不同降质，然后混合训练的扩散模型+增强模型，每一步根据分类的类型处理不同降质。仅仅扩散模型效果可能不好，也可能图片的condition方式不好。微调整个CLIP感觉有点多余，不如用分类器。思想是普通的，实验不太充分。

3. 方法

在这里插入图片描述
我们提出了AutoDIR，这是一种用于图像恢复的自动一体式模型，能够处理多种类型的图像退化，包括弱光、雾等。左：通过AutoDIR进行多任务图像恢复的管道，其中所述盲图像质量评估（BIQA）模块检测损坏图像的主要退化，并用文本提示指示基于潜在扩散的一体式图像恢复模块（AIR）进行恢复。结构校正模块（SCM）进一步改进了恢复图像的细节。对于退化程度未知的图像，例如以绿色虚线为边界的显示不足的相机图像，AutoDIR会自动将任务分解为多个步骤。右图：AutoDIR可以有效地恢复不同退化的干净图像，并可以处理未知退化的图像。（放大查看详细信息）。AutoDIR由一下结构组成：

盲图像质量评估（BIQA）：该模块自动识别输入图像中存在的主要退化，如噪声、模糊、雾度，并生成相应的文本提示，表示为eauto，随后用于图像恢复过程。
一体式图像恢复（AIR）：该模块基于预先训练的潜在扩散模型进行微调。在来自BIQA（或用户提供）的文本嵌入 $e_{auto}$ 的指导下，AIR模块输出恢复的图像 $I_{sd}$ 。
结构校正模块（SCM）：SCM是一个轻量级的卷积网络，旨在进一步恢复图像细节并校正 $I_{sd}$ 中的剩余失真。该恢复过程的输出是最终恢复的图像，表示为 $I_{res}$ 。

3.1 盲图像质量评估（BIQA）

我们使用CLIP模型作为我们的BIQA骨干。然而，由于CLIP是为强调语义信息而非图像质量的视觉识别任务预先训练的，因此它不能直接用于BIQA。下图左可以看到，原始CLIP对于清晰和有雾的图像无法在空间中分离，作者的目的是得到右边的，即不同任务的图像属于特定的空间。
在这里插入图片描述

为了克服这一点，作者分两个步骤来解决这个问题：（i）我们构建了一个新的图像质量评估任务来微调CLIP。（ii）我们为图像质量感知训练提出了一个新的正则化术语。

在这里插入图片描述

首先，作者冻结文本编码器并全微调图像编码器。作者在这里构建了正则本文输入，也就是固定句式只改变 ${ } c_i\{\}$ ，然后将文本输入到 $\epsilon_T$ 中且降质图像输入到 $\epsilon_I$ 中，获得的本文和图像编码计算他们的cosine相似性：

在这里插入图片描述

然后通过softmax取最大概率获得相应的text embedding:

在这里插入图片描述
然后作者使用多类保真度损失(multi-class fidelity loss)对图像编码器进行微调：

在这里插入图片描述

其中， $y(c_i|I)$ 当 $c_i$ 属于 $I$ 图片当降质类型为1，其他为0。这个损失的意思就是最大化 $\hat p(c_i|I)$ 这个概率。这种微调的方式还是无法区分域，如第一张图b所示。作者提出语义不可知损失(semantic-agnostic loss)解决这个问题(见第一张图c)：

在这里插入图片描述

这个损失的意思和 $L_{FID}$ 正好相反，这个损失是降质图像对应的GT匹配当前类型的概率尽可能的小，则最后的总损失函数为：
$L_{BIQA} = L_{FID} + L_{SA}$

3.2 All in one 图像复原 (AIR)

AIR模块是基于潜在扩散模型设计的，利用文本和图像嵌入条件来恢复具有生成先验的图像 $I_{sd}$ 。为了使文本条件能够指导不同的图像恢复，我们结合了交叉注意，将文本条件映射到潜在扩散模型的时间条件UNet主干的中间层。对于图像条件，我们将图像条件 $z_I$ 与噪声潜像 $z_t$ 连接起来，并将它们馈送到UNet主干。

3.3 结构校正模块 (SCM）

具体方法是一个神经网络F，输入是latent diffusion的输出和降质图像concat起来。这么做的话，我们需要先通过扩散采样得到输出 $z$ ，这是费时的。作者直接用任意时间步预测的 $\tilde z$ 来作为扩散采样的输出 $z$ ，损失函数如下：
$L_{SCM} = ||I_{gt}-(F(D(\tilde z), I)+D(\tilde z))||_2^2$
这里F是SCM模型，I是降质图像，D是VAE的解码器，其实F就是学习残差。如下图可以看到，他这个扩散模型的效果很差啊！！！没有这个SCM的话效果一定不行，所以他加SCM这个。