Multi-Stage Progressive Image Restoration

1. 论文

Multi-Stage Progressive Image Restoration（CVPR2021）

• 链接：https://arxiv.org/pdf/2102.02808.pdf
• 代码：http://github.com/swz30/MPRNet

2. 背景

图像恢复任务在恢复图像时，需要在空间细节和上下文信息之间取得复杂的平衡。本文提出了一种协同设计的 MPRNet，可以平衡这些目标。本文主要提出了一个多阶段的架构，逐步学习退化输入的恢复功能，从而将整个恢复过程分解成更容易管理的步骤。

3. 方法

图像恢复框架如图2所示。其包括三个阶段来逐步恢复图像：前两个阶段是基于编码器-解码器（encoder-decoder）的子网络，用来学习较大感受野的上下文信息。由于图像恢复是一个位置敏感的任务，故最后一个阶段采用的子网络在原始输入图像分辨率上进行操作（无降采样），从而在最终的输出图像中保留所需的纹理。

在每两个阶段之间，加入一个监督注意力模块（supervised attention module）。在真实图像的监督下，在将上一阶段的特征图传递给下一阶段之前对其进行重新缩放。此外还引入了跨阶段的特征融合机制（cross-stage feature fusion mechanism），前一个子网络中间的多尺度上下文特征有助于巩固后一个子网络的中间特征。

虽然 MPRNet 堆叠了多个阶段，但每个阶段都能接触到输入图像。我们在输入图像上使用多区域层次结构（multi-patch hierarchy），并将图像分割成非重叠的区域：第一阶段为4个，第二阶段为2个，最后一个阶段为原始图像。

在任何阶段 $S$，本文提出的模型不直接预测恢复后的图像 ${\mathbf{X}}_S$，而是预测一个残差图像 ${\mathbf{R}}_S$，将退化的输入图像 $\mathbf{I}$ 添加到其中，得到 ${\mathbf{X}}_S=\mathbf{I}+{\mathbf{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]$$

其中 $\mathbf{Y}$ 是真值图像， ${\mathcal{L}}_{\text{char}}$ 是 Charbonnier loss，${\mathcal{L}}_{\text{char }}=\sqrt{{\Vert {\mathbf{X}}_S-\mathbf{Y}\Vert }^2+{\epsilon }^2}$，${\mathcal{L}}_{\text{edge }}=\sqrt{{\Vert \Delta \left({\mathbf{X}}_S\right)-\Delta (\mathbf{Y})\Vert }^2+{\epsilon }^2}$，$\Delta$ 是拉普拉斯算子。

3.1. 互补特性处理

目前用于图像修复的CNN通常采用一个编码器-解码器结构，或者一个单尺度的特征流水线。本文提出了一个多阶段的框架，早期阶段结合了编码器-解码器网络，最终阶段采用了一个在原始输入分辨率上运行的网络。

3.1.1 编码器-解码器子网络

图3a为编码器-解码器子网络，该网络基于 U-Net。区别在于：1）增加了信道注意块（channel attention blocks，CAB），以提取每个尺度的特征（CAB见图3b）。2）U-Net跳转连接处的特征图也用CAB处理。3）没有使用转置卷积，而是使用双线性上采样，然后再进行卷积。这有助于减少由于转置卷积而经常出现的输出图像中的棋盘伪影（checkerboard artifacts）。

3.1.2 原始分辨率子网络

为了保留输出图像的细节，模型最后阶段引入了原始分辨率子网络（original-resolution subnetwork，ORSNet）。ORSNet 不采用任何下采样操作。它由多个原始分辨率块（original-resolution blocks，ORB）组成，每个原始分辨率块还包含CAB。ORB的示意图如图3b所示。

3.2. 跨阶段特征融合

本文引入了CSFF（Cross-stage Feature Fusion）模块，用于两个编码器-解码器之间（见图3c）和编码器-解码器与ORSNet之间（见图3d）。优点如下：1）它使网络不容易受到重复使用上采样和下采样操作而导致的信息损失。2）上一阶段的多尺度特征有助于丰富下一阶段的特征。3）网络优化过程变得更加稳定。

3.3. 监督注意力模块

本文在每两个阶段之间引入一个监督注意力模块（supervised attention module，SAM）。SAM的原理图如图4所示。优点如下：1）它提供了对每个阶段渐进式图像恢复有用的真值信号。2）在局部监督预测的帮助下，生成注意力图，抑制当前阶段信息量较小的特征，只让有用的特征传播到下一阶段。

如图4所示，SAM将前一阶段的输入特征 ${\mathbf{F}}_{\mathbf{i}\mathbf{n}}\in {\mathbb{R}}^{H\times W\times C}$ 使用 $1\times 1$ 卷积生成残差图像 ${\mathbf{R}}_S\in {\mathbb{R}}^{H\times W\times 3}$，其中 $H\times W$ 表示空间维度，$C$ 为通道数。将残差图像加到降级后的输入图像 $I$ 上，得到恢复后的图像 ${\mathbf{X}}_S\in {\mathbb{R}}^{H\times W\times 3}$。预测图像 ${\mathbf{X}}_S$ 用真值提供显式监督。接下来，使用 $1\times 1$ 卷积和 sigmoid，从图像 ${\mathbf{X}}_S$ 中生成单像素注意力掩模 $\mathbf{M}\in {\mathbb{R}}^{H\times W\times C}$。然后，这些掩码被用来重新校准变换后的局部特征 ${\mathbf{F}}_{\mathbf{i}\mathbf{n}}$，从而得到注意力引导的特征。最后，由SAM产生的注意力增强特征为 ${\mathbf{F}}_{\mathbf{o}\mathbf{u}\mathbf{t}}$，被传递到下一阶段进行进一步处理。

4. 实验

本文主要进行了去雨（deraining）、去模糊（deblurring）和去噪（denoising）三种实验。

4.1. 去雨结果

4.2. 去模糊结果

4.3. 去噪结果

4.4 消融实验

5. 总结

本文提出了一种用于图像恢复的多阶段体系结构，其中ORB类似SE Block。给人前两阶段使用是U-Net，最后一阶段用的是SE Net的感觉。本文虽然没有特别重大的创新，但工作量比较大，点数也刷的很高。看上去denoising效果似乎稍微差了一点点，说明即使都是low-level的工作，不同的task还是有所区别。