A Unified Conditional Framework for Diffusion-based Image Restoration-Toy模板网

这篇具有很好参考价值的文章主要介绍了A Unified Conditional Framework for Diffusion-based Image Restoration。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

A Unified Conditional Framework for Diffusion-based Image Restoration (Paper reading)

Yi Zhang, CUHK, CN, arXiv2023, Cited:0, Code, Paper

1. 前言

最近，扩散概率模型（Diffusion Probabilistic Models，DPMs）在图像生成任务中表现出了非凡的性能，能够生成高度逼真的图像。当将DPMs用于图像恢复任务时，关键的一点在于如何整合条件信息，以引导DPMs生成准确和自然的输出，这在现有的研究中往往被忽视。在本文中，我们提出了一个基于扩散模型的统一条件框架，用于图像恢复。我们利用一个轻量级的UNet来预测初始引导，并使用扩散模型学习引导的残差部分。通过精心设计扩散模型块的基本模块和集成模块，我们将引导和其他辅助条件信息整合到每个扩散模型块中，实现了空间自适应的生成条件。为了处理高分辨率图像，我们提出了一种简单而有效的逐步分块策略，可以生成任意分辨率的图像而不会产生网格伪影。我们在三个具有挑战性的任务上评估了我们的条件框架：极低光照去噪、去模糊和JPEG恢复，并展示了它在感知质量和泛化到恢复任务方面的显著改进。

2. 整体思想

不算新的思想，用扩散模型对初始估计进行增强。本文创新点应该在扩散模型的网络上，本文是自己设计的网络，其中用到了动态卷积。

3. 方法

A Unified Conditional Framework for Diffusion-based Image Restoration
我们的目标是为图像恢复任务设计一个统一的条件框架。该框架的输入条件信息由两个组成部分组成：降质图像和辅助标量信息。降质图像表示待恢复的图像，而辅助标量信息可以包括降质类型、强度或其他与每个恢复任务相关的细节。

为了增强条件信息的整合，我们首先采用一个轻量级的U-Net来预测初始输出，如图1（左侧）所示。这个初始输出捕捉到了最终恢复图像的低频和确定性方面，这些方面更容易恢复并包含了关键的结构信息。我们将这个初始输出作为扩散模型的空间引导。结合辅助标量信息（例如降级类型、扩散时间步长），我们将它们注入到扩散模型的每个块中，从而实现对扩散模型的更好控制和引导。这种注入不仅提供了全面的上下文，还增强了我们框架的灵活性。我们采用扩散模型来捕捉初始输出的残差分布。

基本模块：在我们的方法中，我们为图像恢复任务中使用的扩散模型设计了一个基本模块。我们的目标是通过利用现有的图像恢复主干网络使模块尽可能简单。我们尽量避免使用复杂的操作符，而是采用现有的图像恢复主干网络，使其尽可能简单。对于每个块，我们使用两个卷积层。在每个卷积层之前，我们引入了LayerNorm来稳定训练过程。我们使用Swish作为激活函数。我们应用了一条捷径来实现残差学习。为了能够注入条件信息，第二个卷积核的设计是基于条件动态变化的。

条件注入模块：为了更好地将条件信息整合到块中，我们提出了一个条件整合模块（Conditional Integration Module，CIM）。在CIM中，首先将引导信息进行缩放，以匹配块内特征图的分辨率。然后，这个缩放后的引导信息通过两个卷积层，并使用SimpleGate激活函数，有效地调整通道数并生成特征图 $G$ 。

SimpleGate(x) = sigmoid(x) * x

同时，辅助标量信息通过两个线性层分支，并在其中使用Swish激活函数进行处理，生成特征图 $S$ 。接下来，特征图 $G$ 和 $S$ 被传递给自适应核引导模块（Adaptive Kernel Guidance Module，AKGM），用于为基本模块中的第二个卷积层生成动态卷积核，如图1所示。AKGM的关键思想是自适应融合卷积核基，使每个空间位置能够根据融合的多源条件信息处理特征图。
A Unified Conditional Framework for Diffusion-based Image Restoration
如图2（左侧）所示，每个AKGM都有 $N$ 个可学习的卷积核基，表示为 $W_b ∈ R^{C×C×k×k}$ ，其中 $C$ 表示通道数， $k$ 表示卷积核大小。这些卷积核基经过训练，用于处理不同的情况和场景。特征图 $G ∈ R^{H×W×N}$ 和 $S ∈ R^{1×1×N}$ 通过逐点乘积融合生成多源融合权重 $M ∈ R^{H×W×N}$ 。这里， $H$ 和 $W$ 分别表示特征图的高度和宽度， $N$ 表示卷积核基的数量。对于特定位置 $(i ， j)$ 的融合卷积核 $F (i ， j)$ ，通过线性融合该位置上的多源融合权重得到。具体而言，可以表示为：
$F_{i,j}=\textstyle \sum_{N-1}^{b=0}M_{i,j}[b]W_b$