基于注意机制的轻量化稠密连接网络单幅图像去雨

柴国强,王大为,芦宾,李竹

(山西师范大学 物理与信息工程学院,临汾 041000)

随着科技的发展,目标检测与跟踪、场景语义分割、深度信息估计[1-3]等计算机视觉算法已经广泛应用于无人机拍摄、人脸识别系统、智慧城市预警等现代化科技中。在阴雨天气下,雨条纹会对相机等传感器获取的图像造成严重的噪声污染,从而降低后续计算机视觉算法的可靠度和准确性,限制了其应用范围。对降雨时获取的图像先进行去雨处理,重构出背景纹理具有重要研究意义。现有图像去雨算法根据所研究的对象不同主要可以分为基于视频的去雨算法和基于单幅图像的去雨算法。基于视频的去雨算法通过分析连续帧图像之间的差异,根据丰富的时空关系检测出雨条纹并将其去除。相比而言,由于缺乏时空信息,单幅图像去雨难度更大。

根据是否需要数据集进行训练,单幅图像去雨算法可以分为基于传统的优化算法和基于深度学习的算法。传统算法中,Li等[4]利用基于高斯混合模型的先验概率对雨条纹进行检测和去除,通过对背景层和雨条纹层施加预先定义的约束实现图像分解,利用结构残基将雨条纹去除。Kim等[5]采用自适应的非局部邻域滤波对雨条纹进行均值滤波将其去除。Luo等[6]提出了一种基于字典学习和稀疏编码的单幅图像去雨算法,其中强互斥性的字典学习逼近雨条纹层和背景层,判别稀疏编码将雨条纹去除。Wang等[7]将字典学习和图像分解相结合实现了单幅图像去雨。

近年来,随着深度学习框架的不断优化和计算机硬件水平的飞速发展,计算机视觉算法性能和运行效率都得到了大幅度提升。Yang等[8]提出了一种基于空洞卷积的多通道单幅图像去雨网络,通过空洞卷积扩大网络感受野,通过闭环反馈实现雨条纹的逐步去除。Wang等[9]先利用空间感知网络得到雨条纹注意图,再利用局部到全局的方式将雨条纹去除。Ren等[10]利用深度特征在不同阶段之间的依赖关系,提出了渐进循环去雨网络。Zhang等[11]提出了一种多任务的去雨网络,在完成去雨的同时实现高分辨率图像的生成。Xia等[12]提出了一种基于图像分解的去雨网络,先将雨图像分解为高频部分和低频部分,再利用简化的稠密连接网络在高频部分将雨条纹去除。王聪等[13-14]分别采用网络与内部尺度融合操作和自校准网络[15]实现了单幅图像去雨。

虽然现有基于深度学习的图像去雨算法取得不错进展,但是就图像去雨质量和网络参数规模方面仍有较大提升空间。本文提出一种基于注意机制的轻量化稠密连接网络,在去除雨条纹的同时降低网络参数规模,提升模型运行效率。

1 网络结构

本文所提出的基于注意机制的改进稠密连接去雨网络结构如图1所示。为了提升雨条纹区域的权重,利用多尺度聚集自适应雨条纹检测模块(multi-scale aggregation self-adapting rain streaks detection module,MASM)检测雨条纹,生成注意图用于引导后续模块有意识地关注降雨区域。雨条纹具有不同的大小和形状,与背景图像相互影响,单一深层网络可能无法从图像中完全去除雨条纹。因此,本文采用循环策略对雨条纹进行分步去除。针对去雨任务,对稠密连接网络进行修改,包括利用空洞卷积增加网络感受野,去除批归一化层(batchnorm layer,BN)增强网络鲁棒性,使用挤压-激励提升网络表征能力。

图1 基于注意机制的改进稠密连接去雨网络结构Fig.1 Structure of modified densely connected rain removal network based on attention mechanism

1.1 多尺度聚集自适应雨条纹检测模块

视觉注意机制有助于计算机视觉算法准确定位于特定区域。为使去雨网络给予雨条纹区域有意识的关注,本文设计了MASM模块检测雨条纹,MASM网络结构如图2所示。

在图2中,输入特征图X经过3个具有不同大小卷积核的并行卷积通道,将生成的特征图沿通道维连接并且经过1×1的卷积核进行特征融合,输出Y表示特征权重,并行卷积通道操作可表示为

图2 多尺度聚集自适应雨条纹检测模块结构Fig.2 Structure of multi-scale aggregation self-adapting rain streaks detection module

式中:X表示输入;Xi(i=1,2,3)表示3个通道的输出,每个通道特征图通道维度为32;F(k)(·)表示卷积核的尺寸为k×k;[X1,X2,X3]表示不同特征图在通道维度上的拼接。

为了检测出小权重雨条纹,首先利用softmax函数对Y进行处理得到归一化的权重B,然后将另一个分支生成的特征图Z与B相乘,最后通过自适应的方式将乘积结果与输入特征X相加:

式中:λ为可学习参数;S为雨条纹注意图。

图3为雨条纹检测模块生成的注意图,可以看出几乎所有雨条纹都得到了标记。

图3 注意图可视化Fig.3 Visualization of attention map

1.2 改进的稠密连接网络

在硬件条件允许的情况下加深网络深度可以有效提升网络性能,但是过深的网络会增加梯度弥散的风险,随着网络深度的增加,前面的信息就会逐渐减弱和消散。为了确保网络中最大化的信息传递,稠密连接网络[16](densely connected networks,DenseNet)使每层都与该层之前的所有层在特征维度上相连,增强了特征的复用,使得信息(前向传播时)或梯度(后向传播时)在整个网络中可以更好的保持,从而训练更深的网络模型。

原始稠密连接网络模块如图4(a)所示,设输入特征图具有c个通道,每个卷积层产生c个通道的新特征图,在网络的前向传递过程中,每一个卷积层都可以接收之前所有卷积层的输出,第n个卷积层的输出可以表示为

图4 稠密连接网络Fig.4 Densely connected network

式中:σ表示激活函数;表示核尺寸为k×k的卷积操作;[X1,X2,…,Xn-1]表示特征图沿通道维度的连接。具有n个卷积层的稠密连接网络模块最终会产生n×G个新的特征图,其中G表示生长率。

在图4(a)所示原始稠密连接网络模块中,每个卷积层后都紧跟一个BN层,在训练过程中BN使用每一批次(mini-batch)的统计量归一化该批次的特征,在测试过程中,BN使用所有训练数据的统计量归一化测试集。在测试集的统计量与训练集差别较大的情况下,BN会影响网络性能,降低网络泛化能力[17]。此外,BN层需要在GPU中存储标准化的数据副本,占用大量计算空间,因此本文将BN层从稠密连接网络模块中去除。

将部分稠密连接块中的常规卷积用空洞卷积(dilated convolution)替代,可以有效增加网络的感受野,在稠密连接网络的最后加入挤压-激励(squeeze-and-excitation,SE)模块,将通道特征和空间特征相结合,进一步增加网络特征提取能力。改进的稠密连接网络如图4(b)所示,其中,输入特征图的通道为64,生长率为12。

1.3 多尺度通道混洗深度可分离卷积

稠密连接网络由于特征复用导致网络后端的卷积核通道维度较大,参数量急剧增加,带来了严重的计算负荷。深度可分离卷积将常规卷积分为深度卷积和逐点卷积,实现了特征图通道维度和空间维度的关联性解耦,可以有效减少网络参数[18-19]。在深度卷积中,每个输出特征层只对单一输入特征层负责,输入特征之间没有信息交换,在逐点卷积中,使用1×1的卷积核实现特征融合。

在上述深度可分离卷积中,不同通道之间的信息交换是通过最后一步卷积核为1×1的卷积实现的,较小的感受野导致了深度可分离卷积不利于特征通道之间的信息融合,限制了网络的表征能力。为了增强不同特征通道之间信息交流,同时获取不同尺度特征,首先将输入特征通道分组,使用不同尺寸卷积核对不同通道进行卷积,然后利用通道混洗技术[20]对不同组的卷积进行特征融合,达到特征通信目的。多尺度通道混洗的深度可分离卷积实现流程如图5所示。

图5 多尺度通道混洗的深度可分离卷积Fig.5 Depthwise separable convolution with multi-scale channel shuffle

2 损失函数

为了确保雨条纹去除的同时恢复出清晰可靠的背景纹理,结合像素目标函数、负结构相似目标函数、感知目标函数、边缘目标函数及注意图目标函数对提出的去雨网络进行训练。

像素目标函数为去雨图像和真实背景RGB幅值的均方差异,其定义为

式中:Lmse(·)为均方误差函数;x为雨图像;G(x)为去雨结果;y为无雨的干净背景。

负结构相似目标函数从亮度、对比度和结构3个方面度量2幅图像的区别,其定义为

式中:SSIM表示结构相似度指数[21],SSIM越大,2幅图像越相近。

感知目标函数[22]从高维度特征差异衡量2幅图像,定义如下:

式中:Vi(·)表示在ImageNet数据集上预训练的VGG-16网络提取的第i层特征图。

采用均方误差训练的网络一般会对生成图像的边缘纹理产生破坏[23],通过增加边缘目标函数来约束网络对边缘纹理的重构,其定义为

式中:W(·)为拉普拉斯边缘检测算子卷积。

注意图目标函数可以有效保证MASM模块能够准确检测出降雨区域,定义为

式中:A(x)为生成的雨条纹注意图;M(x)为雨条纹二值图(0为无雨,1为有雨),M(x)是通过从雨图中减去无雨图像,再通过与预设的阈值比较获取的。

网络总体目标函数为

式中:λ1～λ5为不同目标函数项的权重系数。

3 实验结果和分析

3.1 训练细节

由于无法在实际中获取大量降雨图像与其对应的干净背景图像,采用文献[24]提供的Rain800数据集和文献[25]提供的Rain100H、Rain100L数据集训练和测试相关网络。利用Pytorch框架搭建所提出的神经网络,使用Adam优化器训练模型,初始化学习率为0.01,每30个epoch后学习率降低10倍,共训练100个epoch。在训练过程中,式(9)的权重系数设置为λ1=λ2=λ5=1,λ3=λ4=0.1。采用常用的图像评价指标峰值信噪比(PSNR)和结构相似度(SSIM)对降雨图像进行评价。将本文算法去雨结果与2种基于传统优化的去雨算法DSC[6]、GMM[4]及4种基于深度学习的去雨算法RESCAN[26]、ReHEN[27]、SPANet[9]、UMRL[28]进行比较。

3.2 实验分析

表1为不同算法在测试集上去雨效果定量比较结果。可以看出,本文算法在PSNR和SSIM都取得了最好的表现,并且优势较大,说明了本文算法的有效性。

表1 不同去雨算法PSNR和SSIM对比Table1 PSNR and SSIM comparison of different rain removal algorithms

表2列出了不同算法的参数量及在分辨率为320×640图像上的测试时间。传统优化算法GMM和DSC由于需要学习高斯混合模型或模型字典,非常耗时。基于深度学习的算法由于有预先训练的模型可供使用,运行效率远远高于传统优化算法。得益于多尺度通道混洗深度可分离卷积的利用,本文算法在保证去雨质量的同时大大降低了网络参数,提升了模型推理效率。

表2 不同方法参数量与测试时间对比Table2 Comparison of parameters number and average running time

为了直观地说明本文算法的改进效果,图6给出了2幅合成图像的去雨结果。传统优化算法DSC和GMM很难去除雨纹,背景仍然难以区分。RESCAN和SPANet中产生了较大伪影,ReHEN和UMRL丢失了边缘纹理,细节保持较差。与之对比,本文算法去掉了最多的雨条纹,同时保留了丰富的场景细节,在视觉上更令人满意。

图6 不同算法合成数据集去雨结果可视化对比Fig.6 Visualized comparison of different rain removal agorithms on synthetic dataset

为了验证本文算法的泛化能力,图7展示了不同去雨算法处理真实降雨图像的结果。图像1中雨条纹较小,基于传统算法中的GMM结果使图像产生大范围的模糊,边缘细节保留较差,DSC结果中有残留雨条纹,基于深度学习的算法取得了相近的去雨效果。图像2中雨条纹较大,与其他算法相比,本文算法在去除了几乎所有雨条纹的同时保留了较好的边缘细节。

图7 不同算法真实数据集去雨结果可视化对比Fig.7 Visualized comparison of different rain removal algorithms on real dataset

为了探究阶段数量对网络去雨效果的影响,在Rain100H数据集上采用相同策略训练了包含1～7个阶段的网络。表3列出了不同网络的去雨效果。可以看出,随着去雨阶段从1增加到6,去雨效果改善明显,从6个阶段增加到7个阶段,只有SSIM增加0.1%,达到了网络的瓶颈,为了降低过拟合风险,兼顾去雨质量和运行效率,设定默认包含6个阶段。

表3 阶段数量对去雨效果的影响Table3 Effect of stage number on rain removal

4 结 论

本文提出一种基于轻量化设计的稠密连接网络针对单幅图像进行去雨,从雨条纹检测、雨条纹去除及网络轻量化3方面依次展开,具体贡献如下:

1)利用基于多尺度聚集自适应雨条纹检测模块对雨条纹进行精准定位,生成的注意图用以引导后续模块对雨条纹区域给予有意识的关注;

2)针对图像去雨任务将稠密连接网络进行改进,通过去除归一化层增强网络泛化能力、使用空洞卷积增加网络感受野、利用挤压-激励网络增强表征能力,然后使用循环策略实现雨条纹的去除;

3)提出多尺度通道混洗的深度可分离卷积,替代稠密连接网络中的常规卷积,在保证去雨性能的前提下大幅度降低了网络参数。与现有算法相比,本文算法在合成数据集和真实数据集上都取得了较好的去雨效果,并且运行效率高,验证了算法的效性。