一、主要贡献
•我们提出了DiffIR,一种强大、简单、高效的基于扩散模型的的图像修复方法。与图像生成不同的是,输入图像的大部分像素都是给定的。因此,我们利用DM强大的映射能力来估计一个紧凑的IPR(IR Prior Representation,图像修复的先验表示)来引导图像修复,从而提高DM在图像修复中的恢复效率和稳定性。
•我们建议为Dynamic IRformer提供DMTA(dynamic multi-head transposed attention,动态多头转置注意力机制)和DGFN(dynamic gated feed-forward network,动态门前馈网络),以充分利用IPR。与以往单独优化去噪网络的LDM不同,我们提出联合优化去噪网络和解码器(即DIRformer),进一步提高估计误差的鲁棒性。
•大量实验表明,与其他基于dm的方法相比,提出的DiffIR方法在实现图像修复任务SOTA性能的同时,消耗的计算资源少得多。
二、扩散模型
前向过程:
逆向过程:
目标函数:
三、方法
由于IR中的大部分像素和信息都是给定的,所以对整个图像或feature map进行dm不仅要花费大量的迭代和计算,而且容易产生更多的伪影。为了解决这一问题,我们提出了一种有效的图像重建算法,即DiffIR算法。该算法采用图像重建算法来估计一个紧凑的IPR,从而引导网络进行图像恢复。由于IPR非常小,DiffIR的模型大小和迭代可以大大减少,与传统的DM相比,产生更准确的估计。
注:(1)concat:连接两个数组
(2)pixelunshuffle:一种下采样方法/pixelshuffle:一种上采样方法
(3)GELU:激活函数
(4) F、F_head:输入、输出的feature map
我们将DiffIR的训练分为两个阶段,分别是预训练DiffIR和训练扩散模型。
1.预训练DiffIR
CPEN(compact IR prior extraction network,紧凑图像修复先验提取网络):将GT与低质量图像拼接在一起,经下采样后送入CPEN提取出特征Z即IPR。
DIRformer(dynamic IR transformer,图像修复动态transformer):由DGFN与DMTA组成U-net结构,输入低质量图像,在z的指导下,输出重建的高质量图像。
DGFN(dynamic gated feed-forward network,动态门前馈网络):作用是聚合局部特征。采用1×1 Conv对不同通道的信息进行聚合,采用3 × 3深度Conv对空间相邻像素的信息进行聚合。结构如下:
DMTA(dynamic multi-head transposed attention,动态多头转置注意力机制):作用是聚合全局空间信息,计算attention map,结构如下图
联合训练CPEN和DIRformer,使DIRformer充分利用CPEN提取的图像修复表达IPR,损失函数为:
2.训练扩散模型
利用扩散模型来估计IPR
前向过程:固定CPEN的参数,提取图像修复先验表示Z,对其进行扩散,得到Zt
逆向过程:首先使用CPEN从LQ图像中得到条件向量D,指导去噪网络逐步去除Zt的噪声,经过T次迭代得到估计的Z。
由于IPR是紧凑的,DiffIR可以使用更少的迭代和更小的模型来获得比传统dm更好的估计。由于传统dm在迭代过程中具有巨大的计算成本,因此必须随机采样时间步长t∈[1,t],并仅在该时间步长对去噪网络进行优化。去噪网络与解码器(即DIRformer)缺乏联合训练,意味着去噪网络造成的较小误差,也可能使DIRformer无法发挥其潜力。而DiffIR从第t个时间步长开始,经过所有去噪迭代,得到Zˆ,发送给DIRformer进行联合优化。
损失函数:
四、实验及结论
在inpainting、超分辨率、去运动模糊方面实现了SOTA性能。