Multi-Stage Progressive Image Restoration

🔗 PDF Link 🍺 Github Code

Section 1 介绍

图像恢复是一个从低质量的图像恢复出高清图像的任务。典型的低质量因素有噪声，模糊，水滴噪声，雾噪声等。这是一个高度病态的问题，因为本质上对一张输入的低质量图像存在无限的可行解。为此，许多工作通过增加人工设计的图像先验信息，同时限制解空间的大小为自然图像来解决上述问题。但是设计这种先验本身就已经是一个很具有挑战性的任务了，而且很多时候这种设计的先验的泛化性能不佳。为了改进这个问题，近期的SOTA的方案使用CNNs从大量的数据中学习捕获自然图像的统计先验信息。

基于CNNs的方法的性能表现很大程度上取决于模型结构的设计。大量适用于图像恢复的网络结构模块一集函数单元被设计出来用于提升模型性能，如残差学习，空洞卷积，注意力机制，稠密连接，编解码结构等。然而，几乎所有的这些方法在提出时候(针对低级别的视觉任务)，都基于单阶段设计。与之相反，在高阶的视觉任务中，多阶段的网络模型的显示了更高的效率，例如姿态估计，场景解析以及动作分割。

近期，有一些工作尝试将多阶段的设计引入到去模糊，图像去雨任务中。文章分析了这些方法，找出了限制性能提升的瓶颈点。首先，目前的多阶段方案大都采用了编解码结构，这种方式对大面积的周便信息的编码很有效，但是对于图像空间信息的保留能力缺缺；或使用单尺度的方案，单尺度仅能提供完整的空间信息但是缺少可靠的语义输出。我们尝试将这两种设计整合了起来，变为一个多阶段的结构。其次，我们发现简单地把第一个阶段地结果传递到第二个阶段作为输入本身并不能取得最优的结果(次优)。第三，在图像恢复任务中，给多个阶段提供GT监督非常重要。最后，在多阶段的处理中，在编码解码结构中通过一定的机制将初始的特征图传递给后面的阶段有利于保留周围的纹理特征。

我们提出了一直用多阶段的渐进式图像恢复结构MPRNet，有以下几个主要的组成部分：

1. 前面的stage采用了一个编码解码结构来学习多尺度的纹理信息，而最后的一个stag则在原始的图像尺寸上进行，用以保留完整的空间细节。
2. 一个监督的注意力模块(supervised attention module)被插入到任两个stage之间来促成渐进式的学习。通过与GT的对比，产生的监督信息，SAM模块会尽可能地探索前一个模块的输出来产生注意力图，并根据这个注意力图来实现对前面stage的参数优化，避免传入到下一个stage。
3. 一个跨stage的特征融合机制CSFF(cross-stage feature fusion)被添加以实现多尺度的上下文特征在多个stage之间的传递。同时，这个模块是的多个stages之间的信息流动更便捷了，并是的多阶段网络的优化过程更加稳定。

主要工作贡献:

• 一个新颖的基于多阶段方法，生成具有丰富上下文以及精确的空间信息的能力。基于其多阶段的本质，我们提出的框架将图像恢复的任务拆分为了几个子问题来进行渐进式的恢复。
• 一个有效的监督注意力模块，在多阶段任务中，能在达到下一个阶段之前更有效地利用当前恢复的图像数据。
• 提出了一个跨阶段的多尺度特征融合方法。
• 我们在10个生成数据集以及真实数据集上，展示了MPRNet的性能，取得了SOTA(且在许多任务中有效，复杂度也低，如下图)。同时还做了一些消融实验、量化以及泛化测试。

Section 2 相关工作

早期的图像恢复任务基于total variation(❓)，稀疏编码，自相似度，梯度先验等。近期基于CNN的方法取得了SOTA，从结构设计的角度看，和谐方法大致地可以分为单阶段和多阶段。

Single-Stage Approachs

目前大多数的基于CNN的方法都是单阶段设计，且里面的模块都是为了高层视觉任务所设计。例如残差学习被用于去噪，去雨，去模糊等场景。与之相似，编解码结构，空洞卷积经常被用于提取多尺度的信息。还有一些方法尝试了稠密连接。

Multi-Stage Approaches

多阶段的方法则是尝试采用渐进式的方法，在每一个stage都采用一个轻量级的子网络来实现。这种设计方法将图像恢复的任务分解为了多个简单的子任务来进行，因此会相对更为有效。但是，常规的方法使用同样的子网络并不是一个最好的方案，详见实验4。

Attention

注意力机制被大量地使用与分类、分割、检测等任务中。在图像恢复任务中，例子比比皆是。核心的思想在于注意力机制能够捕捉空间、通道上长距离的相互依赖关系。

Section 3 方法

主要的框架结构如下图所示。

主要包含了三个stages，前两个是基于编解码结构的自网络，利用其较大的感受野用于学习丰富的图像上下文信息。此外，由于图像恢复任务是一个位置敏感的任务(输入输出的像素对保持一致)，最后一个stage则采用了原始尺寸的图像数据作为输入，用以保留想要留下的纹理信息。

相对于直接级联不同的stages，我们采用了一个CAB模块来对接不同的stages。此外，还引入了CSFF机制。前面的stage的中间多尺度语境化特征有助于巩固后一个子网络的的中间层特征。

尽管MPRNet堆叠了多个stages，但是每一个stage都能获取到原始的输入图像数据。与近期的图像恢复方法相似，我们采用了多补丁的层级结构，将数据按照非重叠的形式分块：4个块送入第一个stage，2个块送入第二个stage，原始不分块图像送入到最后一个stage。

在任何一个指定的stage\(S\)，相对于直接预测恢复的图像\(X_S\)，我们提出的模型选择预测残差图像\(R_S\)，因此真实的图像实际为输入图像\(I\)与残差图像的和: \(X_S=I+R_S\)。我们使用如下的损失函数来对MPRNet进行优化： $$ \mathcal{L}=\sum_{S=1}^{3}\left[\mathcal{L}_{\text {char }}\left(\mathbf{X}_{S}, \mathbf{Y}\right)+\lambda \mathcal{L}_{\text {edge }}\left(\mathbf{X}_{S}, \mathbf{Y}\right)\right] $$ 其中\(Y\)表示GT数据，\(\lambda\)被设置为0.05，\(\mathcal{L}_{char}\)表示Charbonnier损失： $$ \mathcal{L}_{\text {char }}=\sqrt{\left|\mathbf{X}_{S}-\mathbf{Y}\right|^{2}+\varepsilon^{2}} $$ 其中\(\varepsilon\)根据经验设置为\(10^{-3}\)。此外，\(\mathcal{L}_{edge}\)表示边缘损失：

$$ \mathcal{L}_{\text {edge }}=\sqrt{\left|\Delta\left(\mathbf{X}_{S}\right)-\Delta(\mathbf{Y})\right|^{2}+\varepsilon^{2}} $$ 其中，\(\Delta\)表示拉普拉斯算子(用于提取边缘)。

后面我 们将会逐个介绍我们方法的关键的组成部分。

3.1 Complementary Feature Processing

目前的用与图像恢复单阶段CNNs方法大都使用如下的一个设计方案：

1. 编解码结构
2. 单尺度特征

Complementary Feature Processing

上图(a)展示了我们的编解码结构子网络，基于标准的UNet，依次添加了一下的一些模块。首先我们添加了CABs用于提取不同不同层次的特征，见上图b；其次，在UNet的跳跃连接部分，我们也加入了CAB；最后，除了使用转置卷积来提升分辨率，我们采用了上线性插值上采样+卷积的形式，这能减少通常由于转置卷积导致的棋盘伪影。

Original Resolution Subnetwork

为了保留精致的细节信息，我们在最后一个stage引入了使用原始图像尺寸的子网络(ORSNet)。最后一个stage没有采用任何的下采样操作，其中包含了多个原始分辨率的blocks(ORBs)，每一个都保留了CAB模块，详见上图b。

3.2 Cross-stage Feature Fusion

在我们的结构中，我们在两个编解码模块之间(上图c)，ORSNet和编解码模块之间(上图d)使用了CSFF模块。注意一个stage之后的特征首先需要经过1x1的卷积之后才会被传递给下一个stage进行整合。CSFF的引入有几个有点。首先，他使得网络更稳定，更不容易因为重复使用上下采样儿导致信息丢失。此外，单stage的多尺度特征能够帮助丰富下一个stage的特征。第三，CSFF的引入使得网络的优化更加稳定，因为其促进的信息的流动。

3.3 Supervised Attention Module

近期用于图像恢复的多阶段网络大都为在每一个stage预测图像，然后将这个结果送入到下一个连续的stage。我们引入了一个监督注意力模块(CAM)来连接两个不同的stage，这个模块的引入显著提升了性能。

CAM的结构如上图所示，该结构的贡献主要分为两步。首先，它提供了一个GT监督信号，这个信号对于渐进式的图像恢复有意义。此外，这个局部监督信号的引入使得只有有效的信息会被送入到下一个stage。

如上图所示，SAM将上一个stage的输入\(\mathbf{F}_{\mathbf{i n}} \in \mathbb{R}^{H \times W \times C}\)经过一个1x1的卷积生成 \(\mathbf{R}_{S} \in \mathbb{R}^{H \times W \times 3}\)。然后将残差图像与低质量的图像相叠加生成恢复图像\(\mathbf{X}_{S} \in \mathbb{R}^{H \times W \times 3}\)。然后使用这个恢复的图像，通过1x1的卷积和一个sigmoid函数生成一些注意力图\(\mathbf{M} \in \mathbb{R}^{H \times W \times C}\)，然后这些注意力图会被用于重构前一个stage的输出特征(SAM模块的输入\(\mathbf{F}_{\mathbf{i n}}\))，生成注意力引导的特征\(\mathbf{F}_{\mathbf{o u t}}\)送入到下一个stage。

Section 4 实验与分析

在多个任务上评估的我们的方法，a)图像去雨，b)图像去模糊，c)图像去噪，共涵盖10个不同的数据集。

4.1 Datasets and Evaluation Protocol

使用PSNR和SSIM来进行对比。后面就不写了，结果自己看表。

Image Deraining

Image Deblurring

Image Denoising

4.2 Implementation Details

我们的MPRNet是端到端训练的且不需要与训练。为不同的任务独立训练模型。我们采用了2CABs在每个尺度的编解码器，使用了2x2和stride为2的max-pooling。在最后一个stage中，采用了包含3个ORBs的ORSNet，其中每个ORB又使用了8个CABs。考虑到任务的复杂度，我们缩放了网络的宽度，去雨用了40个通道，去噪用了80个通道，去模糊用了96个通道。网络都是在256x256的patch上进行训练，batchsize为16，总共训练\(4 \times 10^5\)轮。水平和纵向的反转被随机的应用作为数据增强方法。使用了Adam优化器，初始学习率为\(2 \times 10^{-4}\)，使用cosine annealing策略最终减小到\(1 \times 10^{-6}\)。

4.3 Image Deraining Results

4.4 Image Deblurring Results

4.5 Image Denoising Results

4.6 Ablation Studies

Section 5 高效的图像恢复

资源与效果的平衡是比较难达成平衡的，最理想的方案就是对于资源少的情况有一个低资源需求模型，对于资源多的情况只需要在后面堆叠模块就行。多阶段方案自然而然的就能实现了这个需求。结果如下表所示。

Section 6 结论

阿巴阿巴~