基于注意機制的輕量化稠密連接網絡單幅圖像去雨

柴國強,王大為,蘆賓,李竹

(山西師范大學 物理與信息工程學院,臨汾 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)提出多尺度通道混洗的深度可分離卷積,替代稠密連接網絡中的常規卷積,在保證去雨性能的前提下大幅度降低了網絡參數。與現有算法相比,本文算法在合成數據集和真實數據集上都取得了較好的去雨效果,并且運行效率高,驗證了算法的效性。