基于注意力机制的多尺度特征融合图像去雨方法

0 引言

雨纹会严重影响图像拍摄质量,含有雨纹的图像往往存在图像模糊、图像对比度和能见度降低、目标遮挡等问题,严重阻碍诸如目标检测、目标跟踪和图像识别等计算机视觉任务的执行．因此,图像雨纹的去除研究近年来在计算机视觉领域引起了极大的关注．现有的图像去雨算法包括视频去雨算法和单幅图像去雨算法．视频去雨算法可通过连续图像帧进行雨纹的检测与去除,而单幅图像去雨算法缺乏时序信息,所以单幅图像去雨更具有挑战性．

现有的基于单幅图像的去雨算法可以分为基于图像先验的算法和基于深度学习的算法．Kang等[1]利用双边滤波将图像分解为高频部分和低频部分并在高频部分中应用字典学习去除雨纹．Li等[2]使用高斯混合模型作为先验,将雨纹层从背景图像中分离出来．由于雨纹分布密度不均、形状大小不一,所以基于图像先验的算法无法有效捕捉并去除雨纹．近年来,由于卷积神经网络在计算机视觉领域取得成功,研究人员提出许多基于CNN的图像去雨方法．Fu等[3]提出一种DerainNet网络,利用滤波器将图像分解成高频层与低频层并将高频层送入网络进行训练,随后对网络输出和低频层进行图像增强并将其融合得到最终的去雨图像．Ren等[4]提出一种渐进式去雨网络PReNet,每阶段由残差网络与循环层构成,阶段的结果与原始雨图的拼接作为下一阶段的输入．Fu等[5]设计出一种基于轻量级金字塔去雨网络LPNet,该网络引入拉普拉斯金字塔结合递归残差结构实现图像去雨．Wang等[6]提出一种基于深度跨尺度融合结构的单幅图像去雨网络DCSFN,通过跨尺度方式融合不同的子网并使用内尺度连接方式实现雨纹去除．Yi等[7]提出一种使用残差通道先验引导的去雨网络SPDNet,并通过迭代制导的方式直接生成清晰的无雨图像．

虽然基于深度学习方法的去雨性能方面有了显著改进,但部分网络通过增加深度或者利用多分支结构来增强网络性能,使得结构过于复杂并导致网络难以训练．现有方法并未充分利用多尺度信息,对像素以及通道信息也没有有效利用,去雨后的图像存在局部细节模糊现象且不能很好地保留图像的纹理结构．有鉴于此,本文利用全局特征融合结构可以有效融合不同层次特征的特点,提出一种基于注意力机制的多尺度特征融合图像去雨网络(Multi-scale Feature Fusion Attention Network,MFFAN)．本文的主要工作如下:

1)提出一种端到端的结合注意力机制的多尺度特征融合去雨网络,通过双层特征融合结构加强网络各层间的信息流动与特征复用,在有效避免浅层特征丢失的同时为去雨图像提供更多的细节信息．全局特征融合注意力模块进一步提高了去雨图像质量,且整体网络易于训练．

2)为了充分利用多尺度信息,设计出多尺度注意力残差模块(Multi-scale Attention Residual Block,MARB)．首先利用不同扩张卷积提取不同尺度的特征信息,再通过通道分离—卷积—拼接的方式增强网络的特征提取能力,最后坐标注意力(Coordinate Attention,CA)[8]通过位置信息对通道关系和长期依赖性进行编码,提高网络对于重点特征的关注．

3)使用均方误差以及结构相似度作为联合损失函数约束网络．在仿真和真实雨像数据集上进行评估,所提算法有着较好的去雨纹效果且保留较多的图像细节信息．

1 相关工作

1.1 残差网络

卷积神经网络在计算机视觉领域已经取得巨大成功,学者们提出许多经典的卷积神经网络如LeNet、AlexNet、VGGNet等．通过增加网络的宽度和深度可以提高网络的性能,但是一味地增加深度会导致梯度弥散/爆炸以及网络退化问题．因此,He等[9]提出一种残差网络,并在ImageNet比赛中获得分类和识别任务的冠军．残差网络由多个残差块堆叠而成,在残差块中使用一种跳跃连接,有效缓解了深度神经网络中增加深度所带来的问题．

H(x)=h(x)+F(x),

(1)

其中,x为残差块输入,F(x)为恒等映射输出,F(x)为残差映射．前向传播时,输入信号可以从任意低层直接传播到高层,有效解决网络退化问题;反向传播时,信号可以不经过中间权重层直接到达低层,有效缓解梯度弥散的问题．

1.2 注意力机制

人类可以高效地在复杂场景中找到重要区域．受此启发,注意力机制被引入到计算机视觉中,目的是模仿人类视觉系统．注意力机制在图像分类、目标检测、语义分割等视觉任务中进一步提升了网络的性能．康健等[10]设计出PFA-LinkNet,利用通道注意力机制对编码阶段的特征图各通道间的相互关系进行建模来增强网络提取水体特征的能力．胡序洋等[11]提出一种车道线检测方法LaneSegNet,使用一种混合注意力网络,将空间与通道注意力相结合来获取丰富的车道线特征,增强网络的性能．张凯等[12]提出一种用于射频指纹识别的注意力残差卷积神经网络,通过注意力机制,提高网络对于指纹图像中有效信息的关注．

2 本文方法

2.1 网络结构

本文算法MFFAN的网络结构如图1所示,网络的整体架构包括5个部分:浅层特征提取模块(Shallow Feature Extraction Block,SFEB)、多个串联的多尺度注意力残差组模块(Multi-scale Attention Residual Group Bock,MARGB)、局部特征融合模块(Local Feature Fusion Block,LFFB)、全局特征融合注意力模块(Global Feature Fusion Attention Block,GFFAB)和图像重建模块(Image Reconstruction Block,IRB)．其中每个多尺度注意力残差组由2个多尺度注意力残差模块和局部特征融合模块组成,定义输入图像为xinput．

图1 MFFAN网络框架Fig.1 Network framework of the proposed MFFAN

浅层特征提取模块(SFEB)由一层卷积核大小为3的卷积层构成,对输入图像进行浅层特征提取．

S=Conv3×3(xinput),

(2)

其中,Convi×i表示卷积核大小为i的卷积,S代表浅层提取模块的输出．

多个多尺度注意力残差组模块(MARGB)可表示为

(3)

全局特征融合模块(GFFAB)可表示为

G=fg([M1,M2,M3,…,Mn]),

(4)

其中,fg表示GFFAB,[]表示通道拼接操作,G表示特征融合注意力模块的输出．

局部特征融合模块(LFFB)仅使用通道拼接以及1×1卷积进行融合．最后使用卷积大小为3的卷积核进行图像重建,网络的最终输出可表示为

Y=X-Conv3×3(G),

(5)

其中,Y表示网络的最终输出图像．

2.2 多尺度注意力残差模块

在图像去雨任务中,空间上下文信息对捕捉和去除雨纹十分重要.雨纹的形状和密度都是多样的,使用普通的残差块无法很好地提取到整个图像的雨纹特征．因此,本文使用多尺度注意力残差模块提取雨纹特征,利用多尺度特征提取模块(Multi-scale Feature Extraction Block,MFEB)代替残差块中串联的2个3×3卷积,并使用坐标注意力进一步细化提取特征,具体结构如图2所示．

图2 多尺度注意力残差模块Fig.2 Multi-scale attention residual block

对于模块输入n首先使用2个大小为3,扩张因子分别为1、2的扩张卷积来提取感受野尺度为3×3和5×5的特征信息,并使用LeakyReLu激活函数φ．随后将不同尺度的特征进行拼接,使用1×1的卷积进行通道降维．

k=Conv1×1[φ(Conv3×3,1(n)),φ(Conv3×3,2(n))],

(6)

其中,Convi×i,j表示卷积核大小为i,扩张系数为j的扩张卷积．

为了实现在细粒度层面提取雨纹的多尺度特征,在通道维度使用多级分割卷积拼接方式对雨纹特征k进行提取．首先将k分割成4个具有同等的通道数的特征k1,k2,k3,k4．第1组特征图k1经过卷积核大小为3的卷积,输出特征为o1．将o1在通道维度对半分割成o1,1和o1,2,并将o1,2与下一组特征k2进行通道拼接后送入卷积核大小为3的卷积中．这样操作重复3次,使得拼接后的特征得到充分复用．在前面分组中的特征所经历卷积次数较少,其感受野也较小,这能够帮助网络有效关注细节信息．相反,后面分组中含有的特征经历的卷积次数较多,其感受野也就较大,能够帮助网络关注全局信息．

(7)

最终,将输出特征图o1,1,o2,1,…,o4沿通道维度进行拼接,再使用1×1卷积降低通道维度:

y=Conv1×1([o1,1,o2,1,…,o4])．

(8)

接着使用坐标注意力进一步加强网络性能．坐标注意力首先使用尺寸为(H,1)和(1,W)的池化核沿着输入特征图y的水平坐标及垂直坐标对每个通道使用全局平均池化．因此,高度为h的第c个通道的输出以及宽度为w的第c个通道的输出可以表示为

(9)

(10)

之后将zh与zw在通道维度上进行拼接,并使用一个共享的1×1卷积进行通道降维,紧接着对其使用H-swish非线性激活函数δ,得到中间特征t:

t=δ(Conv1×1(zh,zw))．

(11)

将t沿着空间分割成2个独立分量th和tw,并使用2个1×1的卷积将2个独立分量转换到与输入特征y具有相同的通道数的分量,再使用Sigmoid激活函数,可表示为

mh=σ(Conv1×1(th)),

(12)

mw=σ(Conv1×1(tw)),

(13)

mh和mw分别表示2个坐标轴上坐标的注意力权重,σ表示Sigmoid激活函数．坐标注意力的输出可以表示为

(14)

uc(i,j)表示输出特征图中坐标为(i,j)、通道维度为c的分量．

最终,整体模块的输出M可表示为

M=n+u．

(15)

2.3 全局特征融合注意力模块

简单相加或拼接操作无法很好地融合不同层次的特征．为了加强特征融合的效果,本文使用一种像素注意力为图像中的每个像素赋予不同的权重,使得网络对于雨纹区域的关注度更高．像素注意力[13]应用于图像去雾领域用以解决雾特征在不同像素区域分布不均的问题．首先对不同层次的特征进行通道拼接,并使用一个1×1的卷积进行通道降维,最后送入到像素注意力中,模块整体结构如图3所示.

图3 全局特征融合注意力模块Fig.3 Global feature fusion attention block

全局特征融合注意力模块(GFFAB)可表示为

g=Conv1×1[M1,M2,…,Mi],

(16)

PA=σ(Conv1×1(φ(Conv1×1(g)))),

(17)

G=PA⊗g,

(18)

其中,g表示聚合的全局特征,φ表示ReLu激活函数,PA表示像素注意力模块输出,G表示整体模块输出．

2.4 损失函数

均方误差(Mean Square Error,MSE)作为损失函数被广泛用于图像去雨任务,但会造成重建图像中的高频纹理过度平滑甚至模糊．结构相似度(Structural Similarity,SSIM)损失作为衡量2幅图像之间结构相似性的重要指标,可以使去雨后的图像更符合人类的视觉感官．为了更好地训练网络,将SSIM损失和MSE损失结合形成一种混合损失函数．

(19)

LSSIM=1-SSIM(h(x),I),

(20)

其中,x是网络输入图像,h(x)是网络输出图像,I是无雨标签图像．

通过组合LMSE和LSSIM,最终的混合损失函数可以表示为

L=LMSE+rLSSIM,

(21)

其中,r作为平衡MSE损失与SSIM损失之间的超参数．实验时,将输入图像首先进行归一化处理再送入网络,使得MSE损失与SSIM损失的范围均在[0,1]之间．根据文献[14],将r设置为0.2,在图像保持像素级相似性的同时保持全局结构,获得与原图更接近的去雨图像．

3 实验与结果

3.1 数据集与训练配置

利用Rain100L、Rain100H以及Rain12数据集来对MFFAN进行训练和测试．Rain100L是小雨数据集,由一种类型的雨带的合成图像构成,包含200对训练图像和100对测试图像．Rain100H是大雨数据集,含有5种雨纹方向,包括1 800对训练图像以及100对测试图像．Rain12数据集由于数量较少被作为测试集,将Rain100L作为训练集进行训练．真实世界的雨图可以有效评估算法的泛化性能．

实验时,将每张图像随机裁成大小100×100的图像块进行训练．使用Adam优化器进行模型参数的优化,初始学习率为0.001,训练迭代次数设置为100,在第30、50、80次迭代时,学习率乘以0.2进行衰减,一个批次大小设置为16,卷积核通道数设置为32．实验在一台装有Linux系统,显卡型号为NVIDIA Tesla T4,Python3.7版本,pytorch1.2.0版本的服务器上完成．

3.2 仿真雨像实验与分析

将所提算法在仿真数据集上与JORDER[15]、DRN[16]、PReNet[4]、ReHEN[17]、LPNet[5]和BRN[18]进行对比．图4为本文方法与其他现有方法的去雨图像．表1为各方法在仿真数据集上测试的PSNR(Peak Signal To Noise Ratio,峰值信噪比)、SSIM均值,均在Ycrcb空间取Y通道分量进行计算．在Rain100H、Rain100L数据集上,本文方法获得了最好的结果,PSNR较次优值分别提高0.26 dB和0.45 dB,SSIM分别提高0.5%和0.3%,在Rain12数据集上也获得了较高的评价指标值．

表1 不同算法在仿真雨天数据集上PSNR、SSIM对比

图4 仿真雨天图像多种算法去雨效果对比Fig.4 Comparison of rain removal results by various algorithms on simulated rain images

Rain100H由于雨线密集因此更加具有挑战性．从中选取2张图像进行效果评价,如图4所示:JORDER去除了大部分雨线,但是仍有雨线残留(如上图水草部分,下图天空部分),同时细节恢复不佳;DRN有效去除了雨线,但是放大区域内颜色信息没有有效恢复(如上图鱼眼部分,下图建筑部分);PReNet在放大区域内仍有轻微雨线残留且部分区域模糊不清(如下图放大区域墙壁部分);ReHEN在雨线密集部分处理效果较差(如上图鱼嘴部分),并存在部分细节信息丢失现象(如下图建筑玻璃部分);LPNet去雨效果并不是很好,且纹理十分模糊;BRN过分去除雨纹,导致图像中的白色纹理部分被去除(如上图鱼身部分),同时色彩信息恢复较差(如下图放大区域建筑墙壁部分);MFFAN图像中的雨纹去除较为彻底同时能够保留较多的结构纹理,整体图像与相对应的图像结构更相似．

JORDER利用扩张卷积获得多尺度特征,但忽略了对细粒度层面多尺度特征的有效利用;PReNet、DRN和BRN利用相似的递归结构去除雨纹但却没有充分利用通道信息;ReHEN通过SENet为不同的通道赋予不同的权重,提高了去雨的效率但缺乏对于像素信息的利用;LPNet利用多分支结构获得不同尺度的图像,去雨速度较快但是效果并不理想．相比之下,MFFAN总体网络深度并不深,只有一条主分支避免多分支结构带来的运算成本的提高．MFFAN网络利用多尺度注意力残差模块获得多尺度特征,通过坐标注意力和像素注意力充分利用通道维度和像素维度的信息．简单的双层特征融合结构加强网络各层间流动,且为深度特征提供更多的细节．因此,在3种合成数据集上,MFFAN能够更有效地处理各种形状的雨纹,同时保留图像的细节纹理和色彩信息．

3.3 真实雨像实验与分析

图5为MFFAN与JORDER、DRN、PreNet、ReHEN、LPNet和BRN在真实雨天图像上进行去雨后的结果．JORDER无法彻底去除雨纹在图(第1行第2列)中放大区域内仍有雨纹残留,DRN和BRN有效去除大部分雨纹,但在图(第1行第3列、第2行第3列)左上方区域内产生模糊．PReNet在图(第1行第4列)中放大区域内仍有些许雨纹残留,同时左上方区域细节恢复效果不佳．REHEN无法去除大尺度雨纹,2幅图中放大区域内仍存在雨纹．LPNet去雨后的图像较为平滑,背景纹理细节模糊．而MFFAN展现了较好的泛化性能,去雨后的图像有着很好的视觉效果．

图5 真实雨天图像多种算法去雨效果对比Fig.5 Comparison of rain removal results by various algorithms on real rain images

由于真实世界中的雨像缺乏相应的无雨清晰图像,本文仅从定性方面对现有去雨方法及本文方法进行去雨效果的评估．从分析的结果来看,相较于其他现有的方法,本文所提的MFFAN可以更有效地处理真实环境中的雨纹,获得更符合人类视觉感官的去雨图像．

3.4 消融实验与分析

为了验证所提算法的合理性,在同一环境相同配置下进行消融实验．表2展示了不同数量的残差组数N在Rain100L数据集上的PSNR、SSIM值．

表2 不同数量的残差组在Rain100L上PSNR、SSIM对比

由表2可以看出,当多尺度注意力残差组的数量设置为5时获得了最高的PSNR,但是此时的SSIM值相较于N=4时却没有得到增长．这是由于网络深度达到一定的阈值时,其性能并不能得到有效改善．因此,从运算成本以及去雨性能两方面进行综合考虑,将网络中残差组数量N设置为4．

为了验证本文所提模块的有效性,对不同模块的组合进行消融实验．如表3所示:组合1使用2个串联的3×3卷积代替多尺度特征提取模块并且在模块间进行直连作为基准模型;组合2在多尺度注意力残差模块中使用多尺度特征提取模块;组合3在每两个多尺度注意力残差模块间添加局部特征融合;组合4在图像重建前添加全局特征融合注意力模块;组合5将全局特征融合注意力模块中的像素注意力去除;组合6即本文所提方法．

表3 不同组合模块在Rain100L上PSNR、SSIM对比

从表3中可以明显地看出本文所提网络获得了最高的PSNR和SSIM值．基准模型在增加多尺度特征提取模块后,有效帮助网络提取丰富的雨纹特征．全局特征与局部特征融合增强了网络内部的信息流动性,评价指标得到提高,在多尺度残差注意力模块中添加坐标注意力及在全局特征融合中增加像素注意力即本文所提方法,增强网络的特征提取能力并且得到最高的指标,表明各模块是缺一不可的．

最后在同等环境下验证不同损失函数对去雨效果的影响,结果如表4所示．

表4 不同损失函数在Rain100L上PSNR、SSIM对比

在表4中,可以看到MSE和SSIM的联合损失函数的PSNR比单一损失函数MSE和SSIM分别高出0.4 dB和0.07 dB,SSIM则分别高出0.29%和0.06%．两种指标的提高有效表明了混合损失函数可以获得更好的去雨效果．

消融实验从残差组数量N,网络中不同模块的组合以及损失函数三方面来对网络进行实验分析,将残差组数量N设置为4,使得网络在适当的深度时获得高效的去雨性能,可以避免网络深度过深时,性能却并没有得到提高的问题．本文所提的各模块组合获得了最高的评价指标,说明缺少任意模块都会造成去雨性能的下降,其中多尺度特征提取模块、坐标注意力和像素注意力是本文网络重要的组成模块．最后本文所用的混合损失函数相较于单一的损失函数MSE和SSIM,在Rain100L上取得了最佳的评价指标,表明使用混合损失函数可以获得更好的雨纹处理效果且保留较多的图像细节信息．

4 结论

本文提出一种基于注意力机制的多尺度特征融合图像去雨算法(MFFAN),该算法利用双层特征融合结构将浅层特征进行融合作为后续模块的输入用以充分复用特征,使网络关注局部细节信息的同时关注全局信息．多尺度注意力残差模块充分利用了不同尺度的特征信息,进一步提高网络提取雨纹特征的能力．全局特征融合注意力增强网络对于雨水区域和纹理细节的学习．采用混合损失函数可以获得更好的图像去雨效果．所提算法在Rain100L、Rain100H、Rain12及真实雨像数据集中获得较高的PSNR和SSIM值,在有效去除雨纹的同时很好地保留和恢复图像的细节纹理与色彩信息,泛化性能较好．如何在更加恶劣的天气情况下有效恢复图像以及如何加强模型在真实雨像下的泛化性能将是以后需要改进和解决的问题．