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

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的卷積層構成,對輸入圖像進行淺層特征提?。?/p>

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進行提?。紫葘分割成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值,在有效去除雨紋的同時很好地保留和恢復圖像的細節紋理與色彩信息,泛化性能較好．如何在更加惡劣的天氣情況下有效恢復圖像以及如何加強模型在真實雨像下的泛化性能將是以后需要改進和解決的問題．