【论文阅读】Non-locally Enhanced Encoder-Decoder Network for Single Image De-raining

概述

户外监控设备经常捕捉带有雨纹(rain streak)的图像，这可能会严重影响一些现有电脑视觉系统的表现。因此，自动去除雨纹成为计算机视觉和多媒体领域的重要研究课题，并已成功应用于无人驾驶技术和基于内容的图像编辑等领域。

对于单个图像的去雨处理，现有的基于深度学习的方法要么专注于网络的入口和出口,将输入图像分解为高和低频率信息，使用残差学习减小映射范围，要么引入级联学习(cascaded learning)，将去雨分解为多阶段任务。这些方法将卷积神经网络视为封装的端到端映射模块，没有深入探讨神经网络设计的合理性和优越性。

这篇论文发表于ACM MM 2018，提出一种非局部增强的编码器解码器网络架构，用于增强特征表达和空间相关学习。该架构包含一个嵌入池索引(pooling indices)的编解码器网络，由一系列非局部增强的密集块(dense block)组成，不仅充分利用了所有卷积层的层次特征，而且能够很好地捕获长距离依赖和结构信息，有效学习日益抽象的特征表示。

该方法同时做到有效去除各种密度的rain streak，完美保留图像的细节。
论文链接

---

背景知识

Pooling indices
在Decoder中需要针对最大池化之后的feature map进行上采样。上采样，顾名思义就是池化的反向处理，这个过程中存在着不确定性。比如一个1x1的特征点经过上采样变成一个2x2特征区域，这个区域中的某个1x1区域将会被原来的1x1特征点取代，其他的三个区域为空。但是哪个1x1区域会被原特征点取代呢？一个做法就是随机将这个特征点分配到任意的一个位置，或者干脆给它分配到一个固定的位置。但是这样做无疑会引入一些误差，并且这些误差会传递给下一层。层数越深，误差影响的范围也就越大。所以把1x1特征点放到正确的位置至关重要。
因此通过Pooling Indices的方式来保存池化点的来源信息。在Encoder的池化层处理中，记录最大像素值的位置区域，也就是每一个池化后的1x1特征点来源于之前的2x2的哪个区域，这个信息被称为Pooling Indices。Pooling Indices会在Decoder中使用，利用Decoder对应的池化层的Pooling Indices确定某个1x1特征点应该放到上采样后的2x2区域中的哪个位置。此过程的如下图所示。

（结合知乎两位作者的回答整理而成：
作者：小哼哼
链接：https://zhuanlan.zhihu.com/p/75241386
作者：郭冠华
链接：https://zhuanlan.zhihu.com/p/38871860）

非局部增强
受传统的非局部均值滤波的启发，非局部神经网络被提出，用于实现视频分类任务的远程依赖的计算。在每个二维非局部操作中，一个位置的响应被计算为所有空间位置特征的加权和。
这是神经网络第一个能够将接收域从相邻位置放大到整个图像的组件。
本文中Non-local的使用来自于何凯明博士的《Non-local neural network》，用于处理当前作用点与全局所有信息的关系。
原式如下：

x是输入信号，y是和x尺寸一样的输出信号，i代表时间空间上的输出位置索引，j代表全图中所有可能位置的索引。函数f(x_i, x_j)计算位置i和j的权重。函数g用来计算j位置输入信号的一个表示。C(x)因子用于归一化。
（引自作者：简书程序是我写的。链接：https://www.jianshu.com/p/9991034832ef 来源：简书）
其中函数f有多种选择：Gaussian，Embedded Gaussion，Dot product，Concatenation。

动机

传统的单幅图像去雨方法是探索雨纹物理特性的某些先验信息，将其建模为信号分离问题，对雨纹和无雨场景分别建模；或者直接将其视为图像滤波问题，借助非局部均值平滑(nonlocal mean smoothing)解决。由于这些模型是基于人工设置的底层特征，并固定了一个先验的雨纹假设，它们只能处理特定形状、尺度和密度的雨滴，而且容易导致与雨纹类似的图像细节被破坏。

深度卷积神经网络已广泛应用于单幅图像的去雨处理，并取得了显著的性能改善。这些方法通常将问题建模为像素级的图像回归过程，通过一系列的卷积、池化和非线性操作等，直接利用端到端网络学习将输入的雨纹图像映射到其干净的版本或负残差映射(negative residual map)。
这些模型存在一些局限性：

• 大部分基于深度CNN的模型都是模仿低阶图像处理的经验，如图像去噪、超分辨率和滤波、设计浅层神经网络并在网络传播过程中保持恒定的特征图分辨率。由于网络感知野的大小有限，每个空间位置的像素值推断仅依赖于周围小的局部区域，因此去除较长的雨带通常比较困难。
• 由于对长距离空间上下文建模的忽视，这些模型在检测大雨条纹时，往往难以准确地填充雨滴，删除图像内容，导致结果（特别是纹理丰富的边缘）过于模糊。这些方法都使用CNN作为封装的端到端映射模块，没有深入探讨神经网络设计对于更有效去除雨纹的合理性和优越性。

受自适应非局部均值过滤器的启发，论文建议将非局部操作引入到端到端降雨网络框架的设计中。非局部操作将空间位置上的特征响应计算为某个特征图中特定位置范围内特征的加权和。
提出的框架中，主要有以下重要部分：

• 编码器和解码器都由三个级联的非局部增强密集块(NEDB)组成，以利用雨纹图中丰富的结构线索和无雨自然场景中的自相似性。每个NEDB是一个包含非局部特征映射加权的残差学习模块，然后是四个紧密连接的卷积层和另一个卷积层，前者用于分层特征编码，后者用于残差推断。
• 编码器网络中引入pooling striding机制来学习日益抽象的特征表示，随着接受域的扩大，分辨率降低。
• 加入在相应编码器的最大值池化步骤中计算的pooling indices，以便在解码器中执行非线性上采样，有助于在结果图像中保留结构和细节。

模型

输入输出层

入口处的输入：一幅有雨纹的图像
出口处的输出：无雨图像

Entrance
使用两个卷积层来提取输入雨天图像的浅层特征

$I_0$：输入图像 $H_0$：第一浅层特征提取卷积层
使用残差连接绕过中间层，将输入图像$I_0$和浅层特征$F_0$与接近整个网络出口的层连接起来。
好处在于两方面：

• 提供长距离的信息补偿，使得原始像素值和低级特征激活在整个架构的最后仍然可用
• 残差学习有利于梯度反向传播和像素级预测

接下来将浅层特征$F_0$输入到第二卷积层$H_1$，得到浅层特征$F_1$，作为编码器的输入

Exit
输入图像$I_0$和浅层特征$F_0$加到最终层之前的特征激活上，在最后一层使用tanh得到像素值在(-1,1)的雨纹映射图$R$。最终输出图像可由下式得到

非局部增强编码器和解码器

论文所提出的体系结构可以看作是具有非局部操作和密集连接的增强版本
Non-locally Enhanced Dense Block (NEDB)

NEDB的非局部操作：
在这个工作中使用的该公式如下：

$f$：采用点积形式计算一对值之间的关联：

$F_n$：输入NEDB的特征激活值, 尺寸为$H_n\ast W_n\ast C_n$
$\theta (.)和\varphi (.)$：学习特征，$\theta (F_{n,i})=W_{\theta }{F}_{n,i},\varphi (F_{n,i})=W_{\varphi }{F}_{n,i}$
$g(.)$：计算$F_n$表示的一元函数
$C(F)$：归一化因子 $C(F)={\Sigma }_{\forall j}f(F_{n,i},F_{n,j})$
通过对像素矩阵中每个位置i考虑图像中所有可能位置$(\forall j)$对特征表示进行非局部增强。

将非局部增强过的特征表示作为输入送给五个连续地、紧密相连的卷积层
每一层到所有后续层都直接连接，第$l$层接受之前所有层的特征激活值$D_0,...,D_{l-1}$的拼接作为输入

为避免梯度爆炸/梯度消失问题，在每个NEDB中使用局部残差学习，第m个NEDB的输出为

其中$F(F_{m-1},W_m)$表示将要学习的残差映射，通过对从之前所有层的特征激活值的拼接作1x1卷积得出

Pooling Indices Guided Decoding
编码器部分包含三个连续的NEDB，每个NEDBs后面都有一个带striding的最大池化层，对特征激活值进行下采样。相对应地，在解码器部分堆叠另外三个NEDB，每一个后面都有一个最大池化层对特征激活值进行上采样。利用残差连接将编码层的特征激活与对应的解码层连接。
在编码时为之后解码时的上采样记录pooling indices。根据记录的池化索引矩阵，最大池化层输出的特征映射的计算为：初始化为最大池化操作之前的尺寸，将输入的特征映射上面每个位置的特征值赋给对应的位置，其他位置全部置零。
之后实验将证明该方法比双线性插值上采样方法更适合去雨任务。

Multi-scale Non-Local Enhancement
通过最大池化和最大反池化操作，中间层的特征激活值的尺寸在编码时逐渐减小，在解码时逐渐增加，而NEDB中的非局部操作计算每个空间位置的值之间的相关性，随着尺寸的增大，计算负担显著增加。为了解决这个问题，在构建编码和解码器时，以多尺度的方式实现非局部操作。
具体来说，针对最低空间分辨率的特征激活(如图2中的F4)，后续的NEDB如果直接作用于整个feature activation map，实现全局级的非局部增强，计算量非常大。因此，首先将具有较高空间分辨率的特征激活划分为区域网格(如图2中k×k NEDB表示在执行区域非局部操作之前划分输入特征映射，例如F1被划分为一个8×8的网格），然后让后续的NEDB在每个划分出来的小区域内进行特征激活。同时，与传统的3×3或5×5窗口内的局部卷积运算相比，分区域进行非局部操作能有效检索长距离的结构线索，以便更好地清除雨带。

损失函数

输出图像和ground truth之间的平均绝对误差

$H,W,C,\hat{Y},Y$分别表示高度，宽度，图像的通道数，输出图像和ground truth

实验

数据集

合成数据：

数据集	数量	说明
DDN-Data	14000对,训练9100对，测试4900对	由1000幅清晰的图像合成，包含14种不同的条纹方向和大小
DIDMDN-Data	12000对	3种密度
Rain100L&Rain100H	L：训练200对，测试100对；H：训练1800对，测试100对	L只有1种密度，H有5种方向

真实数据：
网上；《Density-aware Single Image De-raining using a Multi-stream Dense Network》中的数据；自行拍摄
度量标准：
Peak Signal-to-Noise Ratio (PSNR)；Structure Similarity Index (SSIM)

实验过程

• 使用pytorch
• 使用水平翻转增加数据，并调整图像的大小使其长边小于512
• mini batch大小为1，输入图像的大小不必相同
• 优化器：Adam，学习率：0.0005，当训练损失停止减少时减少10%，直到0.0001，weight decay0.0001， momentum0.9

模型对比

定量分析

与五个模型的不同数据集用两个metric进行对比。PSNR显著提高，较高的SSIM值也表明该方法可以更好地恢复图像的结构信息。

定性分析

五幅图像的去除结果的视觉比较如下图。前三个例子是与真实场景相似的合成大雨案例，后两个例子是从Rain100H数据集中选择的较难的例子

现实生活中的图像的去除结果的视觉比较如下图。

Ablation study

对提出的模型NLEDN包含的两个核心组件：带striding和最大反池化操作的编码器-解码器框架，NEDB及非局部特征加权方法，在Rain100L数据集上对模型的五种变形进行有效性和必要性验证。
可以看到$R_c$只是简单地连接多个NEDB可以带来显著的性能改进，这验证了更深层次特征表示在去雨任务中的有效性。实验证明，堆叠6块最好。