DiffBIR

DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior

目标：使用一个统一的框架处理不同的盲图像恢复任务（BSR, BFR, BID）

背景

BSR：StableSR、PASD这些方法需要重新训练

BFR：这些方法只能在人脸图像上有较好性能，泛化性不好

BID：CBDNet、VDNet这些方法虽然能够去除未知噪声，但是结果会过于平滑

zero-shot IR：DDRM、DDNM、GDP这些方法无法完成满意的恢复结果

ControlNet：直接使用原始的带有噪声和伪影的LQ图像会干扰扩散过程，造成不理想的伪影

贡献

1. 将image restoration问题解耦为两个阶段：1）degradation removal：去除退化并获得高保真度的恢复结果；2）information regeneration：生成真实的细节
2. 提出IRControlNet获得真实的图像重建结果
3. 提出 region-adaptive restoration guidance 策略在质量和保真之间折中

方法

动机

LQ图像域庞大复杂，因此包含的条件信息也极为多样化。LQ的退化和内容信息纠缠在一起，因此直接使用LQ作为控制信号会导致不稳定并诱发伪影。

解释： 伪影的生成是由于退化的差异，训练过程中无法从退化图像中区分出内容信息，生成过程受到不可靠条件信息的干扰

因此采用两阶段的方法，生成过程只使用LQ的内容作为条件

Restoration Module

第一阶段只需要去除LQ图像中干扰的退化，不需要生成任何新的内容。由于不同任务之间有不同的特性，因此，“we use separate restoration modules instead of a general one for different BIR tasks to maintain their expertise.”——重要话术

于是，作者直接使用MSE loss训练已有的BIR模型作为Restoration Module：

$$I_{RM}=RM(I_{lq}), \mathcal L_{RM}=\left\|I_{RM}-I_{hq}\right\|_2^2\tag{1}$$

由于要生成多样化的条件图像，因此退化范围设置的比较大

Generation Module

1. 条件编码 使用预训练的固定的VAE编码器 $\mathcal E$ 编码 $I_{RM}$ 到隐空间 ： $c_{RM} = \mathcal E(I_{RM})$
2. 条件网络 使用预训练好的UNet编码器和中间块作为条件网络（记为 $F_{cond}$）。将 $c_{RM}$ 和 $z_t$ concat 作为 $F_{cond}$ 的输入，记为 $ z^{'}t=cat(z_t,c{RM})$ 。

由于通道数增加，因此在 $F_{cond}$ 第一层加入少量参数，并初始化为0，类似于0卷积 ，避免早期训练阶段随机噪声计算梯度

3. 特征调整 类似于ControlNet，只调整中间层特征和跳跃特征

目标函数：

$$\mathcal L_{GM}=\mathbb E_{z_t,c,t,\epsilon, c_{RM}}[\left\|\epsilon-\epsilon_\theta(z_t,c,t,c_{RM})\right\|_2^2]\tag{2}$$

**变体1 **：替换 $\mathcal E$ 为一个小的从0开始训练（trained-from-scratch）的网络

影响：无法保留LQ原始的颜色，且PSNR平均会下降3dB。因此，对于IR任务，条件编码在控制扩散先验中扮演着相当重要的角色。 可能是由于，图像生成过程在隐空间中进行，因此条件需要被投影到相同的空间。

变体2：去掉 $z_t$ ，只将 $I_{RM}$ 作为条件网络的输入

变体3：随机初始化UNet编码器

变体4：控制中间层特征和解码器特征，而不是跳跃特征

从图中可以看出，作者所提出的IRControlNet收敛速度最快。对于变体2，可以得出，$z_t$ 可以加速收敛，可以使得条件网络可以考虑到每个时间步中的随机性，虽然 变体2 PSNR分数较高，但是它的IQA分数较低，往往会生成较平滑的结果，没有足够的纹理细节，总结，$z_t$ 可以加速收敛并帮助生成高质量的结果。对于 变体3，训练损失难以收敛，所有指标中表现最差，因此对条件网络的好的初始化在生成模块中很重要。对于 变体4 其性能与IRControlNet相近，因此控制跳跃特征或解码器特征有相似的表现，然而，解码器的特征通道数是跳跃特征的两倍，因此会引入更多的参数和计算代价。结论：IRControlNet最佳

Restoration Guidance

目的： 设计一个可控的模块完成质量和保真度的权衡

用户往往会期望在高频区域（纹理，边缘）有更多的生成细节，而在平坦区域（天空，墙壁）生成的内容较少

提出区域自适应恢复指导策略，通过用户调节指导scale指导 $\mathcal D(\tilde z_0)$ 面向高保真的 $I_{RM}$ 。使用区域自适应MSE loss，首先使用sobel算子计算梯度，由于一幅图像中有较强梯度幅度的像素很少，因此将 $I_{RM}$ 分为多个不重叠的patchs，计算patch-level的梯度幅度 $\mathcal G(I_{RM})$ ,然后计算权重图 $\mathcal W=1-\mathcal G(I_{RM})$ ：

$$\mathcal L(\tilde z_0) = \frac{1}{HWC}\left\|\mathcal W\odot(\mathcal D(\tilde z_0)-I_{RM})\right\|_2^2\tag{3}$$

梯度较小的区域会赋有较大的权重，反之亦然。于是，低频区域会有较高的loss，会受到高保真的 $I_{RM}$ 的影响更多，而高频区域受到的影响较小，会保留更多生成的内容。

s为指导scale，如：s越大会使得 $\mathcal D(\tilde z_0)$ 更接近于 $I_{RM}$ ，有较高保真度。

实验

参考

[1] Lin X, He J, Chen Z, et al. Diffbir: Towards blind image restoration with generative diffusion prior[J]. arXiv preprint arXiv:2308.15070, 2023.