【图像恢复】Multi-Stage Progressive Image Restoration（译）

仅供参考，如有翻译不到位的地方敬请指出！

多阶段渐进图像恢复

Multi-Stage Progressive Image Restoration

Zamir S W, Arora A, Khan S, et al. Multi-stage progressive image restoration[J]. arXiv preprint arXiv:2102.02808, 2021.

文章下载地址：https://arxiv.org/abs/2102.02808
代码：https://github.com/swz30/MPRNet

摘要

图像恢复任务在恢复图像时需要空间细节和高级上下文信息之间的复杂平衡。在本文中，作者提出了一种新的协同设计，可以最佳地平衡这些竞争目标。主要提出一个多阶段架构，它逐步学习降级输入的恢复功能，从而将整个恢复过程分解为更易于管理的步骤。具体来说，该模型首先使用编码器-解码器架构学习上下文化特征，然后将它们与保留本地信息的高分辨率分支相结合。在每个阶段，我们引入一种新颖的每像素自适应设计，利用原位监督注意力来重新加权局部特征。这种多阶段架构的一个关键要素是不同阶段之间的信息交换。为此，作者提出了一种双向方法，其中信息不仅从早期到后期按顺序交换，而且特征处理块之间还存在横向连接，以避免任何信息丢失。由此产生的紧密互联的多阶段架构称为MPRNet，在包括图像去噪、去模糊和去噪在内的一系列任务中，在十个数据集上提供了强大的性能增益。例如，在Rain100L、GoPro和DND数据集上，与现有技术相比，我们分别获得了4dB、0.81dB和0.21dB的PSNR增益。

1. 引言

图像恢复是从退化版本中恢复清晰图像的任务。退化的典型例子包括噪声、模糊、雨、霾等。这是一个高度不适定的问题，因为存在无穷多个可行解。为了将解空间限制为有效/自然图像，现有的恢复技术明确地使用经验观测人工设计的图像先验。然而，设计这样的先验是一项具有挑战性的任务，而且往往不具有普遍性。为了改善这个问题，最近最先进的方法是采用卷积神经网络(CNNs)，通过从大规模数据中捕获自然图像统计来隐含地学习更一般的先验。
  基于CNN的方法相对于其他方法的性能提升主要归因于其模型设计。已经开发了许多用于图像恢复的网络模块和功能单元，包括递归残差学习，扩张卷积，注意力机制，密集连接，编码器-解码器，以及生成对抗模型。然而，几乎所有这些低级视觉问题的模型都是基于单阶段设计的。相比之下，多阶段网络在高级视觉问题上比单阶段网络更有效，如姿态估计，场景解析和动作分割。
  最近，较少的工作将多阶段设计用于图像去模糊和图像去雨。我们分析这些方法来识别阻碍它们性能的架构瓶颈。首先，现有的多级技术或者采用编码器-解码器架构，该架构在编码广泛的上下文信息方面是有效的，但在保留空间图像细节方面是不可靠的，或者使用单尺度主线，该主线提供空间精确但语义不太可靠的输出。然而，我们表明，在多阶段架构中，两种设计选择的组合对于有效的图像恢复是必需的。其次，我们表明纯粹地将一个阶段的输出传递给下一个阶段会产生次优结果。第三，与[85]不同的是，在逐步恢复的每个阶段提供基本事实监督是很重要的。最后，在多阶段处理期间，需要将中间特征从较早阶段传播到较晚阶段的机制来保留来自编码器-解码器分支的上下文化特征。
  作者提出了一个多阶段渐进图像恢复架构，称为MPRNet，有几个关键组成部分。1).早期阶段使用编码器-解码器来学习多尺度上下文信息，而最后阶段在原始图像分辨率上操作以保留精细的空间细节。2).每两个阶段之间插入一个受监督的注意力模块，以实现渐进学习。在清晰图像的指导下，该模块利用前一阶段的预测来计算注意力图，该注意力图又用于在传递到下一阶段之前细化前一阶段的特征。3).增加了跨阶段特征融合(CSFF)机制，有助于将多尺度上下文化特征从早期传播到后期。此外，该方法简化了各阶段之间的信息流，有效地稳定了多阶段网络优化。
  这项工作的主要贡献是:

• 一个新的多阶段方法，能够产生丰富的上下文和空间准确的输出。由于其多阶段性质，我们的框架将具有挑战性的图像恢复任务分解为子任务，以逐步恢复退化的图像。
• 一个有效的监督注意模块，在进一步传播之前，在细化输入特征的每个阶段充分利用恢复的图像。跨阶段聚合多尺度特征的策略。
• 通过在10个合成和真实世界数据集上设置新的最先进的技术来证明MPRNet的有效性，这些数据集用于各种恢复任务，包括图像去噪、去模糊和去噪，同时保持较低的复杂性。此外，提供了详细的消融、定性结果和一般测试。

2. 相关工作

近年来，高端DSLR相机向智能手机相机转变。然而，用智能手机摄像头捕捉高质量图像是一项挑战。由于相机的限制和/或不利的环境条件，图像中经常出现图像退化。早期恢复方法基于全变差，稀疏编码，自相似，梯度先验等。最近，基于CNN的恢复方法取得了最新的成果。就结构设计而言，这些方法可以大致分为单阶段和多阶段。
  单阶段方法：目前，大多数图像恢复方法基于单阶段设计，并且架构组件通常基于为高级视觉任务开发的组件。例如，残差学习已经用于执行图像去噪，图像去模糊和图像去噪。类似地，为了提取多尺度信息，通常使用编码器-解码器和扩张卷积模型。其他单级方法包含密集连接。
  多阶段方法：这些方法旨在通过在每一阶段采用轻量级子网以渐进方式恢复清晰的图像。这样的设计是有效的，因为它将具有挑战性的图像恢复任务分解成更小更容易的子任务。然而，一个常见的做法是在每个阶段使用相同的子网，这可能会产生次优的结果，如第4部分的实验所示。
  注意力机制：在图像分类、分割和检测等高级任务的成功驱动下，注意力模块已经被用于低级视觉任务。例子比比皆是，包括图像去噪，去模糊，超分辨率，去噪。其主要思想是沿着空间维度、通道维度或两者捕捉相互依赖特征。

【以下几个多阶段图像恢复的文章可以结合着看，能够帮助更好的理解多阶段网络的思想】

1. Nah S, Hyun Kim T, Mu Lee K. Deep multi-scale convolutional neural network for dynamic scene deblurring[C]. Proceedings of the IEEE conference on computer vision and pattern recognition. 2017:3883-3891.
2. Li X, Wu J, Lin Z, et al. Recurrent squeeze-and-excitation context aggregation net for single image deraining[C]. Proceedings of the European Conference on Computer Vision (ECCV). 2018: 254-269.
3. Tao X, Gao H, Shen X, et al. Scale-recurrent network for deep image deblurring[C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018: 8174-8182.
4. Ren D, Zuo W, Hu Q, et al. Progressive image deraining networks: A better and simpler baseline[C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2019: 3937-3946.
5. Fu X, Liang B, Huang Y, et al. Lightweight pyramid networks for image deraining[J]. IEEE transactions on neural networks and learning systems, 2019, 31(6): 1794-1807.
6. Zheng Y, Yu X, Liu M, et al. Residual Multiscale Based Single Image Deraining[C]. British Machine Vision Conference(BMVC). 2019: 147.
7. Zamir S W, Arora A, Khan S, et al. Cycleisp: Real image restoration via improved data synthesis[C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2020: 2696-2705.
8. Suin M, Purohit K, Rajagopalan A N. Spatially-attentive patch-hierarchical network for adaptive motion deblurring[C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2020: 3606-3615.

3. 论文提出的方法

如图2所示，所提出的图像恢复框架包括三个逐步恢复图像的阶段。前两个阶段是基于编码器-解码器子网络，由于大的感受野，这些子网络学习广泛的上下文信息。由于图像恢复是一个位置敏感的任务(需要从输入到输出的像素到像素的对应关系)，最后一个阶段采用了一个在原始输入图像分辨率上操作的子网络(没有任何下采样操作)，从而在最终输出图像中保留所需的精细纹理。
            图2 渐进图像恢复的多阶段体系结构
  本模型不是简单地级联多个阶段，而是在每两个阶段之间加入一个受监督的注意力模块。在清晰图像的监督下，该模块在将前一阶段的特征图传递到下一阶段之前对其进行重新缩放。此外，引入了跨阶段特征融合机制，其中早期子网络的中间多尺度上下文化特征有助于巩固后一子网络的中间特征。
  虽然MPRNet堆叠了多个阶段，但每个阶段都可以访问输入图像。类似于最近的恢复方法，我们在输入图像上调整多面片层次结构，并将图像分割成不重叠的面片:四个用于阶段-1，两个用于阶段-2，原始图像用于最后阶段，如图2所示。
  在任何给定的阶段S，不是直接预测恢复的图像$X_S$，所提出的模型预测残差图像$R_S$，退化的输入图像 $I$ 被添加到该残差图像$R_S$以获得:$X_S=I+R_S$。我们通过以下损失函数优化MPRNet端到端:
$$L=\sum _{S=1}^3[L_{char}(X_S,Y)+\lambda L_{edge}(X_S,Y)](1)$$
其中 $Y$ 代表真实图像，而$L_{char}$代表Charbonnier损失:
$$L_{char}=\sqrt{{\Vert X_S-Y\Vert }^2+{\epsilon }^2}(2)$$
对于所有实验，常数 $\epsilon$ 根据经验设置为$10^{-3}$。此外，边缘损耗$L_{edge}$定义为:
$$L_{edge}=\sqrt{{\Vert \Delta (X_S)-\Delta (Y)\Vert }^2+{\epsilon }^2}(3)$$
其中 $\Delta$ 表示拉普拉斯算子。等式(1)中的参数 $\lambda$ 控制两个损失项的相对贡献，与文献[36]中同样设置为0.05。接下来，描述提出方法的每个关键部分。

3.1 互补特征处理

现有的用于图像恢复的单级CNNs通常使用以下架构设计之一：1).编-解码器，2).单尺度特征主线。编-解码器网络首先逐渐将输入映射到低分辨率表示，然后逐渐应用反向映射来恢复原始分辨率。虽然这些模型有效地编码了多尺度信息，但是由于重复使用下采样操作，它们易于牺牲空间细节。相比之下，在单尺度特征主线上运行的方法在生成具有精细空间细节的图像方面是可靠的。然而，由于有限的感受野，它们的输出在语义上不太健壮。这表明了上述架构设计选择的固有局限性，这些设计选择能够产生空间上准确或上下文上可靠的输出，但不能同时产生两者。为了利用这两种设计的优点，我们提出了一个多级框架，其中早期阶段结合了编-解码器网络，最后阶段采用了一个在原始输入分辨率上运行的网络。
  编-解码器子网：图3a显示了编-解码器子网，它基于标准的U-Net，具有以下组件。首先，添加通道注意块(CABs) 来提取每个尺度的特征(CABs见图3b)。第二，U-Net跳过连接的特征映射也用CAB处理。最后，我们不是使用转置卷积来提高解码器中特征的空间分辨率，而是使用双线性上采样，然后是卷积层。这有助于减少输出图像中由于反卷积而经常出现的棋盘格伪影。

图3 （a）编-解码器子网，（b）ORSNet子网中原始解析块的说明，（c）阶段1和阶段2之间的跨阶段特征融合，（d）第二阶段和最后阶段之间的CSFF。
  原始分辨率子网：为了保留从输入图像到输出图像的精细细节，模型在最后阶段引入了原始分辨率子网(ORSNet)(见图2)。ORSNet不采用任何下采样操作，并生成空间丰富的高分辨率要素。它由多个原始解决方案块(ORB)组成，每个解决方案块还包含CABs。ORB的示意图如图3b所示。

3.2 跨阶段特征融合

在本文的框架中，在两个编-解码器之间(见图3c)以及编-解码器和ORSNet之间(见图3d)引入了CSFF模块。请注意，在将一个阶段的特征传播到下一个阶段进行聚合之前，先用1×1卷积对其进行细化。提出的CSFF有几个优点。首先，由于在编-解码器中重复使用上采样和下采样操作，使得网络不容易受到信息丢失的影响。第二，一个阶段的多尺度特征有助于丰富下一个阶段的特征。第三，网络优化过程变得更加稳定，因为它简化了信息流，从而允许我们在整个体系结构中添加几个阶段。

3.3 监督注意模块

最近用于图像恢复的多阶段网络直接预测每个阶段的图像，然后将其传递到下一个连续阶段。相反，我们在每两个阶段之间引入一个受监督的注意模块，这有助于实现显著的性能增益。SAM的示意图如图4所示，其贡献是双重的。首先，它在每个阶段提供了对渐进图像恢复有用的基本事实监控信号。其次，在局部监督预测的帮助下，我们生成注意力图来抑制当前阶段信息量较少的特征，只允许有用的特征传播到下一阶段。

              图4 监督注意模块
  如图4所示，SAM在较早阶段获取输入特征 $F_{in}\in R^{H\times W\times C}$，并首先通过简单的1 × 1卷积生成残差图像 $R_S\in R^{H\times W\times 3}$，其中 $H\times W$ 表示空间维度，C表示通道数量。将残差图像加到退化的输入图像I上，得到复原图像 $X_S\in R^{H\times W\times 3}$。对于这个预测的图像 $X_S$，我们提供了真实图像的明确监督。接下来，使用1×1卷积随后sigmoid激活从图像 $X_S$ 生成的像素注意力掩模 $M\in R^{H\times W\times C}$。然后，这些掩模被用于重新校准变换后的局部特征 $F_{in}$ (在1×1卷积之后获得)，导致注意力引导的特征被添加到身份映射路径。最后，由SAM产生的注意力增强的特征表示 $F_{out}$ 被传递到下一阶段进行进一步处理。

4. 实验分析

针对几个图像恢复任务，包括(a)图像去噪，(b)图像去模糊，以及©10个不同数据集上的图像去噪，对我们的方法进行了评估。

4.1 数据集和评估协议

使用PSNR和SSIM 指标进行定量比较。
            表1 各种图像恢复任务的数据集描述

图像去雨：使用最近图像去雨的最佳方法的相同实验设置，我们在从多个数据集收集的13712个清晰-有雨图像对上训练我们的模型，如表1所示。利用这个单一的训练模型，我们对各种测试集进行评估，包括Rain100H 、Rain100L 、Test100、Test2800 和Test1200 。
图像去模糊：正如在[68，85，41，69]中一样，我们使用GoPro数据集，该数据集包含2103个图像对用于训练，1111个图像对用于评估。此外，为了证明通用性，我们采用我们的GoPro训练模型，并将其直接应用于HIDE 和RealBlury数据集的测试图像。HIDE数据集是专门为人类感知的运动去模糊而收集的，其测试集包含2025幅图像。当合成生成GoPro和HIDE数据集时，真实模糊数据集的图像对是在真实世界条件下捕获的。真实模糊数据集有两个子集:(1)RealBlur-J由相机JPEG输出形成，以及(2)RealBlur-R通过对原始图像应用白平衡、去马赛克和去噪操作离线生成。
图像去噪：为了训练模型图像去噪任务，我们使用了SIDD数据集的320幅高分辨率图像。对来自SIDD数据集的1280个验证patches和来自DND基准数据集的1000个patches进行评估。这些测试patches是由原始作者从全分辨率图像中提取的。SIDD和DND数据集都由真实图像组成。

4.2 实现细节

MPRNet是端到端可训练的，不需要预处理。我们使用以下设置为三个不同的任务训练单独的模型。我们在编-解码器子网络的每个尺度上使用2个CABs，对于下采样，使用2×2最大池，步长为2。在最后一个阶段，采用包含3个ORBs的ORSNet，每个ORBs进一步使用8个CABs。该网络在256×256个块上训练，批量大小为16，迭代次数为$4\times 10^5$次。为了增加数据，水平和垂直翻转是随机应用的。我们使用初始学习率为$2\times 10^{-4}$的Adam优化器，并使用余弦退火策略，该学习率稳步降低到$1\times 10^{-6}$。

4.3 图像去雨结果

对于图像去雨任务，与先前的工作一致，我们使用Y通道(在YCbCr颜色空间中)计算图像质量分数。表2显示，我们的方法通过在所有五个数据集上持续获得更好的PSNR/SSIM分数，显著提高了技术水平。与最近的最佳算法MSPFN 相比，我们获得了1.98dB的性能增益(所有数据集的平均值)，表明误差减少了20%。一些数据集的改进高达4 dB，例如Rain100L。此外，我们的模型比MSPFN少3.7倍的参数，同时快2.4倍。

      表2 图像去雨结果，最佳和分数第二的会突出显示并加下划线。

  图5显示了挑战性图像的视觉对比。我们的MPRNet在消除不同方向和大小的雨条纹方面是有效的，并且生成视觉上令人愉快并且忠实于基本事实的图像。相比之下，其他方法会损害结构内容(第一行)，引入人为因素(第二行)，并且不能完全消除雨条纹(第三行)。
              图5 图像去雨结果

4.4 图像去模糊结果

我们在表3中报告了在合成GoPro 和HIDE数据集上评估的图像去模糊方法的性能。总体而言，我们的模型相对于其他算法表现良好。与以前的最佳表现技术相比，我们的方法在PSNR数据集上实现了9%的改进，在SSIM数据集上实现了21%的改进，在HIDE数据集上实现了11%和13%的误差减少。值得注意的是，我们的网络仅在GoPro数据集上进行训练，但在HIDE数据集上获得了最先进的结果(+0.98分贝)，从而证明了其强大的泛化能力。

            表3 图像模糊结果（GoPro和HIDE数据集上）。

  我们在两个实验设置下，在最近的真实模糊数据集的真实世界图像上评估我们的多平面重建网:1).将GoPro训练的模型直接应用于RealBlur(测试对真实图像的泛化)，以及2).对真实模糊数据进行训练和测试。表4显示了实验结果。对于设置1，与DMPHN算法相比，我们的MPRNet在RealBlur-R子集上获得了0.29dB的性能增益，在RealBlur-J子集上获得了0.28dB的性能增益。在设置2中观察到类似的趋势，其中我们在SRN上的增益分别为0.66dB和0.38dB。

         表4 RealBlur数据集上的去模糊比较结果。

  图6显示了通过评估方法得到的一些去模糊图像。总的来说，我们的模型恢复的图像比其他模型更清晰，更接近真实。
              图6 图像去模糊结果

4.5 图像去噪结果

在表5中，我们展示了PSNR/SSIM在SIDD和DND 数据集上对几种图像去噪方法的评分。我们的方法比最先进的方法获得了相当大的增益，即，在SIDD上比Cycliesp高0.19dB，在DND上比SADNet高0.21dB。请注意，DND数据集不包含任何训练图像，即完整的公开发布的数据集只是一个测试集。在DND基准上用我们的SIDD训练模型进行的实验结果表明模型可以很好地推广到不同的图像领域。

         表5 图像去噪（SIDD和DND数据集上）。

  图7示出了视觉结果。我们的方法能够去除真实的噪声，同时保留图像的结构和纹理细节。相比之下，通过其他方法恢复的图像要么包含过于平滑的内容，要么包含带有斑点纹理的伪像。

              图7 图像去噪结果

4.6 消融实验

在这里，我们进行消融实验，以分析我们的模型的每个组成部分的贡献。对GoPro数据集进行评估，在大小为128×128的图像块上训练去模糊模型$10^5$次迭代，结果如表6所示。

       表6 对MPRNet的单个部件进行消融研究

级数：随着阶段数量的增加，我们的模型产生更好的性能，这验证了我们的多阶段设计的有效性。
子网的选择：由于我们模型的每个阶段可能采用不同的子网设计，我们测试不同的选项。我们表明，与对所有级采用相同的设计相比，在早期使用编-解码器，在最后一级使用ORSNet，可提高性能(29.7dB)(U-Net+UNet为29.4dB，ORSNet+ORSNet为29.53dB)。
SAM和CSFF：我们通过从最终模型中移除所提出的监督注意模块和跨阶段特征融合机制来证明它们的有效性。表6显示了PSNR从30.49dB到30.07dB的大幅度下降，当我们去掉CSFF时从30.49dB到30.31dB。移除这两个组件会大幅降低性能，从30.49dB降至29.86dB。

4.7 资源高效的图像恢复

CNN模型通常表现出准确性和计算效率之间的权衡。在追求更高精度的过程中，通常会开发更深、更复杂的模型。尽管大型模型往往比小型模型表现更好，但计算成本可能会高得令人望而却步。因此，开发资源高效的图像恢复模型是非常有意义的。一种解决方案是通过每次目标系统改变时调整其容量来训练相同的网络。然而，这是繁琐的，往往不可行。一种更理想的方法是拥有一个单一的网络，该网络能够(a)对计算有效的系统进行早期预测，以及(b)进行后期预测以获得高精度。多阶段恢复模型自然提供了这样的功能。
  表7展示了我们的多阶段方法的阶段结果。MPRNet在每个阶段都表现出了竞争性的恢复性能。值得注意的是，第一阶段模型是轻的，快速的，并产生比其他复杂的算法，如SRN 和de blockanv 2]更好的结果。类似地，与最近的方法DMPHN 相比，我们的第二阶段模型显示PSNR增益为0.51dB，同时更具资源效率(参数减少2倍，速度加快13倍)。

         表7 MPRNet的逐级去模糊性能

5. 结论

在这项工作中，我们提出了一个多阶段的图像恢复架构，通过在每个阶段注入监督来逐步改善退化的输入。我们为我们的设计制定指导原则，要求在多个阶段进行互补的特征处理，并在它们之间进行灵活的信息交换。为此，我们提出了上下文丰富和空间精确的阶段，它们统一编码一组不同的特征。为了确保交互阶段之间的协同作用，我们提出了跨阶段的特征融合和从早期到后期的注意力引导输出交换。我们的模型在众多基准数据集上实现了显著的性能提升。此外，我们的模型在模型大小方面是轻量级的，在运行时间方面是高效的，这对于资源有限的设备来说是非常有趣的。