基于联合注意力机制的单幅图像去雨算法

徐成霞，阎庆，李腾，2，苗开超

（1.安徽大学电气工程与自动化学院，合肥 230601；2.计算智能与信号处理教育部重点实验室（安徽大学），合肥 230601；3.安徽省气象局安徽省公共气象服务中心，合肥 230091）

0 引言

近年来，深度学习在计算机视觉中被广泛应用，在图像识别［1］、图像去噪［2］等领域取得巨大成功。Fu 等［3］结合深度残差网络［4］，建立无雨图像和雨图的负向残差映射，缩小映射范围，优化映射空间，并通过实验验证了算法的去雨能力。Zhang 等［5］搭建一种自动确定雨线密度信息的多流密集网络，根据雨线密度评估出相应的标签，根据标签信息指导去雨任务。算法在有标签的场景中取得良好的效果，但对于现实给出的标签存在偏差。Wang 等［6］构建含有29.5×103对真实降雨数据集，提出基于空间注意力的图像去雨网络（Spatial Attentive Network，SPANet），算法利用空间注意力从局部到全局定位雨线信息，重构无雨背景。然而单一的空间注意力忽略了特征通道的信息不一致性，无法有效抑制冗余信息，突出关键特征。

基于此，本文提出通道注意力和空间注意力的级联模式，利用通道注意力对不同通道赋予不同权重，实现对特征通道的抑制或突出，同时联合空间注意力对特征通道簇进行差异性的特征提取。其次，使用改进的深度残差收缩模块替换原有的标准残差模块，通过软阈值子网络的自动学习以逐渐抑制数据集中的噪声和冗余信息，提升算法对不同降雨场景的适用性。实验结果表明，相较于主流去雨算法，本文算法的去雨效果更佳。此外，本文构建了多场景的降雨数据集（Multiple Scenes Rain，MSR），补充降雨场景的多样性。

1 构建MSR

数据集是深度学习取得成功的重要因素，目前去雨算法大多为有监督的训练方式，输入与监督信号为有雨/无雨图像对。Yang 等［7］构建的Rain100L&R100H，训练集有2 000对，测试集有200 对。Zhang 等［8］构建的Rain800 中训练集700 对，测试集100 对。Zhang 等［5］提出更大规模的合成数据集DID-MDN（Density-aware single Image De-raining using a Multi-stream Dense Network），数据集包含3 种密度的雨线，共12 000 对。合成场景添加的噪声存在规律性较强、雨线分布单一、噪声痕迹明显等问题，去雨算法在训练时容易对相同噪声产生依赖性，不适用于规律较为复杂的真实场景。

Wang 等［6］提出半自动的算法，结合雨线的时间特性和人为监督，在静态背景的降雨序列中获取无雨图像，为去雨任务提供了很大的帮助；但其场景不丰富，图像的视野较小，对于数据驱动的深度学习算法，不能有效地训练去雨算法归纳大场景的降雨规律，导致雨线残留。针对上述问题，本文构建了MSR，如图1 所示。

图1 MSR雨图示例Fig.1 Examples of rain image in MSR

MSR 合成数据集分为三步：首先生成不同密度的随机噪声以模拟不同大小的雨量；其次将噪声拉长、旋转方向，模拟雨线的多样性；最后将生成的雨滴噪声和原始图像叠加，合成雨图。雨线密度为轻度、中度、重度；雨线倾斜角度方面，以垂直方向为0°，分别有向右倾斜5°、10°、15°、20°、25°、30°，向左倾斜5°、10°、15°、20°、25°、30°，以及部分角度的混合叠加，共计20 种不同方向的合成雨线。合成降雨数据集共有6 456 张无雨图片，32 280 张合成雨图，图片的分辨率为256×256。MSR 多样性的合成雨线，为算法提升去雨精度提供数据保障。此外，MSR 通过引用文献［6］中的半自动算法，在大量静态背景降雨序列中获取真实降雨数据集，共11 437对，图像分辨率为256×256。不同的是，MSR 真实降雨数据集大多采用视野宽阔的场景，并且不对数据集进行过度切割，避免造成场景视野较小，无法有效帮助去雨算法在较大场景中的精确去雨的问题。

2 联合注意力去雨算法

2.1 深度残差收缩网络

图像去雨任务中，雨图经常含有雨线之外的噪声，例如光晕和雾气，会模糊图像细节，破坏场景的背景重构，造成较为严重的滤波效应，因此有效地抑制此类噪声成为提升去雨算法性能的途径之一。

软阈值为信号降噪算法常用技术，它将小于设定阈值的信号置0。而软阈值的设置需要通过不断调试，缺乏扩展性。Zhao 等［9］提出将软阈值学习与残差模块结合，形成残差收缩模块，如图2（a）所示。网络在高噪声的振动信号特征学习中，自适应调整阈值，有效地抑制噪声信息，提升网络的特征学习能力。

图2（b）为软阈值学习模块，首先对残差单元最后一层的输出进行绝对值化，经全局平均池化（Global Average Pooling，GAP）获取与上一层相同卷积核数量的一维向量，通过2 层全连接网络和激活函数，将其参数缩放至（0，1）内，其公式为：

图2 残差模块结构Fig.2 Residual module structure

其中：zc为第二层全连接网络的第c个神经元。输出αc为相应的缩放参数，其阈值如下：

其中：τc表示特征图第c个通道的阈值；w、h分别为特征图的宽和高。将缩放参数α与|x|的均值相乘，得到阈值。其目的是将软阈值控制在合理范围内，避免输出特征全为零。原深度残差收缩网络适用于一维信号，本文对其进行改进，使其应用于图像数据。后续实验结果表明，深度残差收缩模块的加入，提升了网络对于光晕、雾气等噪声抑制作用，有效地缓解此类噪声对于背景重构的影响。

2.2 联合注意力去雨网络

2.2.1 联合注意力机制

去雨算法需要检测雨线位置信息，定位雨水需要扩大感受野，获取上下文信息。目前用于定位雨水的模块主要有扩张网络［10］、非局部模块［11］和注意力机制［12］，本文采用联合注意力机制定位雨线。注意力机制可以提升网络筛选信息的效率，抑制不相关信息。对于主干网络的特征图F∈RC×H×W，其中，C、H、W依次为特征图通道数、高度、宽度。如式（3）和式（4）所示，注意力机制将特征F通过通道注意力机制Mc∈RC×1×1得到F′，将F′通过空间注意力机制Ms∈R1×H×W得到F"，⊗为矩阵逐元素相乘，如图3（a）所示。

通道注意力将特征图在空间维度上进行压缩，获取一维矢量，如图3（b）所示。其中最大池化（Max Pooling，MaxPool）和平均池化（Average Pooling，AvgPool）用于聚合特征映射的空间信息，将信息输入到共享网络，压缩输入特征的空间维度，得到，逐元素求和，产生通道注意力图：

其中：σ为激活函数；MLP为共享网络；W0、W1分别为共享网络的第一层和第二层。

空间注意力如图3（c），将通道注意力模块的特征图作为输入，在通道维度分别进行了平均值池化和最大值池化，得到，将它们进行通道叠加，再依次经过卷积、激活函数，产生空间注意力图：

图3 注意力模块Fig.3 Attention mechanism modules

其中：f7×7代表卷积核尺寸7×7 的卷积层。对于中间特征图，会沿着两个独立的维度依次推断注意力图，然后将注意力图与输入的特征图相乘以进行自适应特征优化。

2.2.2 联合注意力去雨网络

联合注意力模型（Joint Attentive Model，JAM）利用通道注意力对不同特征通道中雨线信息的检测，赋予特征通道不同的权重，表征某通道与关键信息的关联度。收集各通道与雨线的关联度后，使用空间注意力以局部至全局的方式检测每个特征通道的雨线及邻域的位置信息，进行雨线的详细定位，如图4（d）。JAM 在主干网络中采用四向循环神经网络（Recurrent Neural Networks with ReLU and Identity matrix initialization，IRNN）层［14］，IRNN 通过叠加2 个循环神经网络（Recurrent Neural Network，RNN）层［15］积累大范围多变的信息。通过信息在整张图中进行有效的传播，改善无规律的真实雨线的检测效果。

图4 联合注意力去雨网络架构Fig.4 Joint attention de-raining network architecture

本文在深度残差收缩模块的基础结构上嵌入联合注意力结构，构成联合注意力残差收缩模块（Joint Attentive Residual Shrinkage Block，JARSB）。在抑制光晕和雾气等噪声的前提下，利用联合注意力图聚焦重要特征，提升网络对于信息的筛选性能。联合注意力单元（Joint Attentive Unit，JAU）由JAM 和JARSB 组成，其中JAM 生成的注意力图，引导JARSB 通过负向映射消除雨线。在联合注意力网络（JOint Attention Network，JOANet）中，网络重复使用4 次JAU 结构，来自于JAM 的注意力权重共享，最后通过两个深度残差收缩模块和一层卷积重构去雨后的背景。

3 实验与结果分析

3.1 实验环境和评价指标

本文在公开数据集Rain100H&Rain100L［7］、DID-MDN［5］、SPANet-Data［6］和本文合成的MSR 上进行实验，并与目前主流单幅图像去雨算法进行比较。

实验采用结构相似性（Structural SIMilarity，SSIM）［16］、峰值信噪比（Peak Signal-to-Noise Ratio，PSNR）［17］、基于像素的视觉信息保真度（Pixel-based Visual Information Fidelity，VIFP）［18］作为评价指标。算法使用PyTorch 实现，并在4 块NVIDIA 1080Ti 显卡上运行。算法训练采用Adam 优化器，单个训练批次的样本数量根据GPU 的显存上限设为16。初始学习率设为10-5，训练终止条件为组合损失停止下降。

3.2 实验结果分析

3.2.1 注意力图的对比与分析

实验将验证注意力机制改进的效果，实验采用SPANet-Data 为训练数据集，对由SPANet 和本文网络生成的注意力图进行对比。图5（a）表示某位置为噪点的概率图，深蓝色到深红色代表该点像素为雨滴的可能性由低变高。

由图5（c）、（d）中红色方框区域可知：图5（d）的红色方框检测的雨线密度大于图5（c）中红色方框的雨线密度，佐证联合注意力在检测雨线方面优于SPANet 的空间注意力。表明通过对每个特征通道包含雨线信息的分析，进行自适应的雨线检测，可以更有针对性、全面地检测雨线的位置信息，指导网络完成去雨任务。

图5 空间注意力与联合注意力的雨线检测结果Fig.5 Rain streak detection results for spatial attention and joint attention

3.2.2 去雨算法对比实验及分析

将本文算法与其他主流去雨算法进行对比，合成数据集随机选用了Rain100H 和DID-MDN 中50%的数据集，真实降雨数据集采用SPANet-Data；并且为了验证实验中去雨算法在各种场景中去雨的鲁棒性，测试数据集由合成降雨数据集Rain100H、DID-MDN 和真实降雨数据集SPANet-Data 组成，表1 中给出了各算法详细的去雨结果的评价指标。如表1 所示，基于联合注意力机制的单幅图像去雨算法在DID-MDN和SPANet-Data 上均有优异的表现，尤其在SPANet-Data 中，3 种评价指标均是最高；但在Rain100H 的效果一般，原因在于本文网络的优势在于使用联合注意力机制检测雨线，而Rain100H 中，合成雨线的痕迹较重，严重破坏了背景中的细节特征，导致重构图像的质量一般。

表1 不同去雨算法结果对比Tab.1 Comparison of results of different de-raining algorithms

不同去雨算法的可视化结果如图6 所示。

图6 不同去雨算法的可视化结果Fig.6 Visualization results of different de-raining algorithms

3.2.3 MSR验证实验

对于数据驱动的去雨算法，训练数据集的质量对结果会产生不可忽略的影响。此部分为验证数据集对去雨算法效果的影响，采用MSR-Synthetic、Rain100H&Rain100L、DIDMDN 训练本文提出的去雨算法，测试结果如表2 所示。

表2 合成降雨数据集上的交叉验证结果Tab.2 Cross-validation results on synthetic rain data

从表2 中发现本文数据集MSR-Synthetic 在Rain100H 上的客观评价指标低于Rain100H&Rain100L 的去雨算法。原因有二：

其一，MSR-Synthetic 并没有Rain100H 同类型的噪声，虽然雨线的检测和去除部分良好，但是存在去雨重构图像过于平滑问题；

其二，Rain100H 中人工添加的雨线噪声过于密集，严重影响了图片背景信息的表达，细节信息明显不足，对比图7的Rain100L 的测试数据集可验证，同样的去雨算法和背景，对于噪声相对较轻的雨图，去雨重构后保留的细节也相对较多。

图7（b）、（c）、（d）分别为JOANet 使用Rain100H&100L、DID-MDN、MSR-Synthetic 作为训练数据集获取的去雨算法。交叉验证中，MSR-Synthetic 在Rain100L、DID-MDN、MSRSynthetic 合成数据集中测试评价指标相对较高。因为这些降雨数据集的噪声分布较为均匀，噪声程度较轻，在原始雨图上保留了较多的背景信息。其中，在DID-MDN 上的测试结果优于DID-MDN 作为训练集得到的去雨算法的测试结果。因为MSR-Synthetic 的噪声结构与DID-MDN 类似，且数据集种类与数量要优于后者，所以MSR-Synthetic 在DIDMDN 上的测试评价指标高于DID-MDN 的去雨算法在本身数据集上的测试指标。

图7 合成降雨数据集上的去雨结果Fig.7 De-raining results on synthetic rain data

本文对真实降雨数据集也做了交叉验证，结果如表3 所示。MSR-Real 的去雨算法在SPANet-Data 和MSR-Real 中的测试指标均有优异的表现，尤其在SPANet-Data 数据量远大于MSR-Real 的情况下，在MSR-Real 中的测试指标依旧高于SPANet-Data 去雨算法的测试指标。原因主要为SPANet-Data 与MSR 真实降雨数据集的风格不一，前者将一帧图片分割为若干子图，若输入尺寸维持在256×256 时，雨滴相对变大，而MSR 采用的训练数据集为完整的一帧图片，有利于训练网络在真实环境检测雨线的能力。可视化效果见图8。

图8 真实降雨数据集上的去雨结果Fig.8 De-raining results on real rain data

表3 真实降雨数据集上的交叉验证结果Tab.3 Cross-validation results on real rain data

3.2.4 残差收缩模块降噪实验

针对部分去雨任务中光晕等噪声破坏场景的背景重构，导致滤波效应的问题，本文去雨算法采用残差收缩模块替换原有的标准残差块。为了客观评价两种残差块影响去雨算法的程度，实验中基础架构采用SPANet，将SPANet 中的标准残差块替换为残差收缩模块。训练数据集采用SPANet-Data，测试数据集来源于网络数据。

对比图9（a）、（b）分析发现：图（b）重构图像存在局部细节模糊，主要原因为该场景中存在强光源，造成雨线周围的细节信息被干扰，不易提取，从而对重构背景的效果造成影响。图9（c）图像在保持原有图像细节尽量不丢失的条件下，实现图像去噪。由此验证残差收缩模块有助于提升算法在学习过程中对于图像噪声的抑制作用。

图9 残差收缩模块降噪实验Fig.9 Denoising experiment of residual shrinkage module

3.2.5 消融实验

本文分析了网络各模块对于去雨算法的影响，如表4 所示。其中：Ba代表空间注意力模块（Spatial Attentive Module，SAM）与标准残差块（Residual Module，RB）的组合，Bb代表SAM 与残差收缩模块（Residual Shrinkage Module，RSB）的组合，Bc代表通道注意力模块（Channel Attentive Module，CAM）与RSB 的组合，Bd代表SAM、CAM 与RSB 的组合。

表4 网络模块分析Tab.4 Network module analysis

训练和测试的数据集均为MSR 真实降雨数据集，测试样本共1 000 对图像对。结合表4，由Ba和Bb可知，深度残差收缩块相较于标准残差块，其PSNR 有小幅提升，说明深度残差收缩块在标准残差块进行深度特征提取的基础上，进行了噪声抑制，消除冗余信息，一定程度上加强了背景重构部分的细节特征。结合Bd分支代表的网络结构在对比实验中的表现，佐证本文提出的基于联合注意力的去雨算法，不仅可以利用联合注意力精准地定位雨线信息，而且可以改善雨图重构的细节特征。Bd相较于Ba在测试数据集上的表现，PSNR 提升4.5%，SSIM 提升0.3%，验证了本文算法提高图像去雨的可行性。

此外，本文在消融实验环节添加运行时间复杂度分析，测试环境不变。结合表4，针对不同分辨率在不同算法架构中的运行时间可知，基础框架Ba在两种分辨率雨图中的运行时间分别为0.103 s 和0.341 s，将标准残差块替换后，运行时间分别上升了22.33%、24.05%，其本质是因为深度残差收缩块在标准残差块的基础上增加软阈值学习模块抑制噪声信息，由此增加网络的学习成本。对比Bb、Bc可知，添加通道注意力机制的学习成本低于空间注意力机制的学习成本，其原因是两个注意力机制属于不同维度：通道注意力属于通道维度，它更关注特征通道间的相关性，基于特定的任务，学习不同通道的重要性；空间注意力属于尺度空间维度，关注的是单个特征通道内不同区域对于任务的贡献度，其粒度比通道注意力细致，故学习成本较高。从整体观测，JOANet 在不同分辨率中的运行时间相比基础框架Ba分别增加了28.16%、29.33%，相较于算法在去雨精度上的提升，在时间学习成本上还需要进一步优化。

4 结语

本文针对单幅图像去雨，提出了基于联合注意力的卷积神经网络，关联不同维度的注意力机制，提出利用通道注意力捕捉各通道雨线的密集程度；同时，利用空间注意力以局部到全局的方式积累上下文信息，实现精准去雨。此外，本文引入深度残差收缩模块并对其改进，降低部分图片中非雨线类型噪声对于背景重构的影响。根据公开数据集以及本文构建的数据集中的实验结果显示，本文算法的去雨性能优于其他算法，但本文去雨算法存在一定局限性，学习时间成本需要进一步优化，这也是今后需要改善的方向。