基于注意機制的輕量化稠密連接網(wǎng)絡(luò)單幅圖像去雨

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

(山西師范大學(xué) 物理與信息工程學(xué)院,臨汾 041000)

隨著科技的發(fā)展,目標(biāo)檢測與跟蹤、場景語義分割、深度信息估計[1-3]等計算機視覺算法已經(jīng)廣泛應(yīng)用于無人機拍攝、人臉識別系統(tǒng)、智慧城市預(yù)警等現(xiàn)代化科技中。在陰雨天氣下,雨條紋會對相機等傳感器獲取的圖像造成嚴(yán)重的噪聲污染,從而降低后續(xù)計算機視覺算法的可靠度和準(zhǔn)確性,限制了其應(yīng)用范圍。對降雨時獲取的圖像先進(jìn)行去雨處理,重構(gòu)出背景紋理具有重要研究意義?，F(xiàn)有圖像去雨算法根據(jù)所研究的對象不同主要可以分為基于視頻的去雨算法和基于單幅圖像的去雨算法?；谝曨l的去雨算法通過分析連續(xù)幀圖像之間的差異,根據(jù)豐富的時空關(guān)系檢測出雨條紋并將其去除。相比而言,由于缺乏時空信息,單幅圖像去雨難度更大。

根據(jù)是否需要數(shù)據(jù)集進(jìn)行訓(xùn)練,單幅圖像去雨算法可以分為基于傳統(tǒng)的優(yōu)化算法和基于深度學(xué)習(xí)的算法。傳統(tǒng)算法中,Li等[4]利用基于高斯混合模型的先驗概率對雨條紋進(jìn)行檢測和去除,通過對背景層和雨條紋層施加預(yù)先定義的約束實現(xiàn)圖像分解,利用結(jié)構(gòu)殘基將雨條紋去除。Kim等[5]采用自適應(yīng)的非局部鄰域濾波對雨條紋進(jìn)行均值濾波將其去除。Luo等[6]提出了一種基于字典學(xué)習(xí)和稀疏編碼的單幅圖像去雨算法,其中強互斥性的字典學(xué)習(xí)逼近雨條紋層和背景層,判別稀疏編碼將雨條紋去除。Wang等[7]將字典學(xué)習(xí)和圖像分解相結(jié)合實現(xiàn)了單幅圖像去雨。

近年來,隨著深度學(xué)習(xí)框架的不斷優(yōu)化和計算機硬件水平的飛速發(fā)展,計算機視覺算法性能和運行效率都得到了大幅度提升。Yang等[8]提出了一種基于空洞卷積的多通道單幅圖像去雨網(wǎng)絡(luò),通過空洞卷積擴(kuò)大網(wǎng)絡(luò)感受野,通過閉環(huán)反饋實現(xiàn)雨條紋的逐步去除。Wang等[9]先利用空間感知網(wǎng)絡(luò)得到雨條紋注意圖,再利用局部到全局的方式將雨條紋去除。Ren等[10]利用深度特征在不同階段之間的依賴關(guān)系,提出了漸進(jìn)循環(huán)去雨網(wǎng)絡(luò)。Zhang等[11]提出了一種多任務(wù)的去雨網(wǎng)絡(luò),在完成去雨的同時實現(xiàn)高分辨率圖像的生成。Xia等[12]提出了一種基于圖像分解的去雨網(wǎng)絡(luò),先將雨圖像分解為高頻部分和低頻部分,再利用簡化的稠密連接網(wǎng)絡(luò)在高頻部分將雨條紋去除。王聰?shù)萚13-14]分別采用網(wǎng)絡(luò)與內(nèi)部尺度融合操作和自校準(zhǔn)網(wǎng)絡(luò)[15]實現(xiàn)了單幅圖像去雨。

雖然現(xiàn)有基于深度學(xué)習(xí)的圖像去雨算法取得不錯進(jìn)展,但是就圖像去雨質(zhì)量和網(wǎng)絡(luò)參數(shù)規(guī)模方面仍有較大提升空間。本文提出一種基于注意機制的輕量化稠密連接網(wǎng)絡(luò),在去除雨條紋的同時降低網(wǎng)絡(luò)參數(shù)規(guī)模,提升模型運行效率。

1 網(wǎng)絡(luò)結(jié)構(gòu)

本文所提出的基于注意機制的改進(jìn)稠密連接去雨網(wǎng)絡(luò)結(jié)構(gòu)如圖1所示。為了提升雨條紋區(qū)域的權(quán)重,利用多尺度聚集自適應(yīng)雨條紋檢測模塊(multi-scale aggregation self-adapting rain streaks detection module,MASM)檢測雨條紋,生成注意圖用于引導(dǎo)后續(xù)模塊有意識地關(guān)注降雨區(qū)域。雨條紋具有不同的大小和形狀,與背景圖像相互影響,單一深層網(wǎng)絡(luò)可能無法從圖像中完全去除雨條紋。因此,本文采用循環(huán)策略對雨條紋進(jìn)行分步去除。針對去雨任務(wù),對稠密連接網(wǎng)絡(luò)進(jìn)行修改,包括利用空洞卷積增加網(wǎng)絡(luò)感受野,去除批歸一化層(batchnorm layer,BN)增強網(wǎng)絡(luò)魯棒性,使用擠壓-激勵提升網(wǎng)絡(luò)表征能力。

圖1 基于注意機制的改進(jìn)稠密連接去雨網(wǎng)絡(luò)結(jié)構(gòu)Fig.1 Structure of modified densely connected rain removal network based on attention mechanism

1.1 多尺度聚集自適應(yīng)雨條紋檢測模塊

視覺注意機制有助于計算機視覺算法準(zhǔn)確定位于特定區(qū)域。為使去雨網(wǎng)絡(luò)給予雨條紋區(qū)域有意識的關(guān)注,本文設(shè)計了MASM模塊檢測雨條紋,MASM網(wǎng)絡(luò)結(jié)構(gòu)如圖2所示。

在圖2中,輸入特征圖X經(jīng)過3個具有不同大小卷積核的并行卷積通道,將生成的特征圖沿通道維連接并且經(jīng)過1×1的卷積核進(jìn)行特征融合,輸出Y表示特征權(quán)重,并行卷積通道操作可表示為

圖2 多尺度聚集自適應(yīng)雨條紋檢測模塊結(jié)構(gòu)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]表示不同特征圖在通道維度上的拼接。

為了檢測出小權(quán)重雨條紋,首先利用softmax函數(shù)對Y進(jìn)行處理得到歸一化的權(quán)重B,然后將另一個分支生成的特征圖Z與B相乘,最后通過自適應(yīng)的方式將乘積結(jié)果與輸入特征X相加:

式中:λ為可學(xué)習(xí)參數(shù);S為雨條紋注意圖。

圖3為雨條紋檢測模塊生成的注意圖,可以看出幾乎所有雨條紋都得到了標(biāo)記。

圖3 注意圖可視化Fig.3 Visualization of attention map

1.2 改進(jìn)的稠密連接網(wǎng)絡(luò)

在硬件條件允許的情況下加深網(wǎng)絡(luò)深度可以有效提升網(wǎng)絡(luò)性能,但是過深的網(wǎng)絡(luò)會增加梯度彌散的風(fēng)險,隨著網(wǎng)絡(luò)深度的增加,前面的信息就會逐漸減弱和消散。為了確保網(wǎng)絡(luò)中最大化的信息傳遞,稠密連接網(wǎng)絡(luò)[16](densely connected networks,DenseNet)使每層都與該層之前的所有層在特征維度上相連,增強了特征的復(fù)用,使得信息(前向傳播時)或梯度(后向傳播時)在整個網(wǎng)絡(luò)中可以更好的保持,從而訓(xùn)練更深的網(wǎng)絡(luò)模型。

原始稠密連接網(wǎng)絡(luò)模塊如圖4(a)所示,設(shè)輸入特征圖具有c個通道,每個卷積層產(chǎn)生c個通道的新特征圖,在網(wǎng)絡(luò)的前向傳遞過程中,每一個卷積層都可以接收之前所有卷積層的輸出,第n個卷積層的輸出可以表示為

圖4 稠密連接網(wǎng)絡(luò)Fig.4 Densely connected network

式中:σ表示激活函數(shù);表示核尺寸為k×k的卷積操作;[X1,X2,…,Xn-1]表示特征圖沿通道維度的連接。具有n個卷積層的稠密連接網(wǎng)絡(luò)模塊最終會產(chǎn)生n×G個新的特征圖,其中G表示生長率。

在圖4(a)所示原始稠密連接網(wǎng)絡(luò)模塊中,每個卷積層后都緊跟一個BN層,在訓(xùn)練過程中BN使用每一批次(mini-batch)的統(tǒng)計量歸一化該批次的特征,在測試過程中,BN使用所有訓(xùn)練數(shù)據(jù)的統(tǒng)計量歸一化測試集。在測試集的統(tǒng)計量與訓(xùn)練集差別較大的情況下,BN會影響網(wǎng)絡(luò)性能,降低網(wǎng)絡(luò)泛化能力[17]。此外,BN層需要在GPU中存儲標(biāo)準(zhǔn)化的數(shù)據(jù)副本,占用大量計算空間,因此本文將BN層從稠密連接網(wǎng)絡(luò)模塊中去除。

將部分稠密連接塊中的常規(guī)卷積用空洞卷積(dilated convolution)替代,可以有效增加網(wǎng)絡(luò)的感受野,在稠密連接網(wǎng)絡(luò)的最后加入擠壓-激勵(squeeze-and-excitation,SE)模塊,將通道特征和空間特征相結(jié)合,進(jìn)一步增加網(wǎng)絡(luò)特征提取能力。改進(jìn)的稠密連接網(wǎng)絡(luò)如圖4(b)所示,其中,輸入特征圖的通道為64,生長率為12。

1.3 多尺度通道混洗深度可分離卷積

稠密連接網(wǎng)絡(luò)由于特征復(fù)用導(dǎo)致網(wǎng)絡(luò)后端的卷積核通道維度較大,參數(shù)量急劇增加,帶來了嚴(yán)重的計算負(fù)荷。深度可分離卷積將常規(guī)卷積分為深度卷積和逐點卷積,實現(xiàn)了特征圖通道維度和空間維度的關(guān)聯(lián)性解耦,可以有效減少網(wǎng)絡(luò)參數(shù)[18-19]。在深度卷積中,每個輸出特征層只對單一輸入特征層負(fù)責(zé),輸入特征之間沒有信息交換,在逐點卷積中,使用1×1的卷積核實現(xiàn)特征融合。

在上述深度可分離卷積中,不同通道之間的信息交換是通過最后一步卷積核為1×1的卷積實現(xiàn)的,較小的感受野導(dǎo)致了深度可分離卷積不利于特征通道之間的信息融合,限制了網(wǎng)絡(luò)的表征能力。為了增強不同特征通道之間信息交流,同時獲取不同尺度特征,首先將輸入特征通道分組,使用不同尺寸卷積核對不同通道進(jìn)行卷積,然后利用通道混洗技術(shù)[20]對不同組的卷積進(jìn)行特征融合,達(dá)到特征通信目的。多尺度通道混洗的深度可分離卷積實現(xiàn)流程如圖5所示。

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

2 損失函數(shù)

為了確保雨條紋去除的同時恢復(fù)出清晰可靠的背景紋理,結(jié)合像素目標(biāo)函數(shù)、負(fù)結(jié)構(gòu)相似目標(biāo)函數(shù)、感知目標(biāo)函數(shù)、邊緣目標(biāo)函數(shù)及注意圖目標(biāo)函數(shù)對提出的去雨網(wǎng)絡(luò)進(jìn)行訓(xùn)練。

像素目標(biāo)函數(shù)為去雨圖像和真實背景RGB幅值的均方差異,其定義為

式中:Lmse(·)為均方誤差函數(shù);x為雨圖像;G(x)為去雨結(jié)果;y為無雨的干凈背景。

負(fù)結(jié)構(gòu)相似目標(biāo)函數(shù)從亮度、對比度和結(jié)構(gòu)3個方面度量2幅圖像的區(qū)別,其定義為

式中:SSIM表示結(jié)構(gòu)相似度指數(shù)[21],SSIM越大,2幅圖像越相近。

感知目標(biāo)函數(shù)[22]從高維度特征差異衡量2幅圖像,定義如下:

式中:Vi(·)表示在ImageNet數(shù)據(jù)集上預(yù)訓(xùn)練的VGG-16網(wǎng)絡(luò)提取的第i層特征圖。

采用均方誤差訓(xùn)練的網(wǎng)絡(luò)一般會對生成圖像的邊緣紋理產(chǎn)生破壞[23],通過增加邊緣目標(biāo)函數(shù)來約束網(wǎng)絡(luò)對邊緣紋理的重構(gòu),其定義為

式中:W(·)為拉普拉斯邊緣檢測算子卷積。

注意圖目標(biāo)函數(shù)可以有效保證MASM模塊能夠準(zhǔn)確檢測出降雨區(qū)域,定義為

式中:A(x)為生成的雨條紋注意圖;M(x)為雨條紋二值圖(0為無雨,1為有雨),M(x)是通過從雨圖中減去無雨圖像,再通過與預(yù)設(shè)的閾值比較獲取的。

網(wǎng)絡(luò)總體目標(biāo)函數(shù)為

式中:λ1～λ5為不同目標(biāo)函數(shù)項的權(quán)重系數(shù)。

3 實驗結(jié)果和分析

3.1 訓(xùn)練細(xì)節(jié)

由于無法在實際中獲取大量降雨圖像與其對應(yīng)的干凈背景圖像,采用文獻(xiàn)[24]提供的Rain800數(shù)據(jù)集和文獻(xiàn)[25]提供的Rain100H、Rain100L數(shù)據(jù)集訓(xùn)練和測試相關(guān)網(wǎng)絡(luò)。利用Pytorch框架搭建所提出的神經(jīng)網(wǎng)絡(luò),使用Adam優(yōu)化器訓(xùn)練模型,初始化學(xué)習(xí)率為0.01,每30個epoch后學(xué)習(xí)率降低10倍,共訓(xùn)練100個epoch。在訓(xùn)練過程中,式(9)的權(quán)重系數(shù)設(shè)置為λ1=λ2=λ5=1,λ3=λ4=0.1。采用常用的圖像評價指標(biāo)峰值信噪比(PSNR)和結(jié)構(gòu)相似度(SSIM)對降雨圖像進(jìn)行評價。將本文算法去雨結(jié)果與2種基于傳統(tǒng)優(yōu)化的去雨算法DSC[6]、GMM[4]及4種基于深度學(xué)習(xí)的去雨算法RESCAN[26]、ReHEN[27]、SPANet[9]、UMRL[28]進(jìn)行比較。

3.2 實驗分析

表1為不同算法在測試集上去雨效果定量比較結(jié)果。可以看出,本文算法在PSNR和SSIM都取得了最好的表現(xiàn),并且優(yōu)勢較大,說明了本文算法的有效性。

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

表2列出了不同算法的參數(shù)量及在分辨率為320×640圖像上的測試時間。傳統(tǒng)優(yōu)化算法GMM和DSC由于需要學(xué)習(xí)高斯混合模型或模型字典,非常耗時?；谏疃葘W(xué)習(xí)的算法由于有預(yù)先訓(xùn)練的模型可供使用,運行效率遠(yuǎn)遠(yuǎn)高于傳統(tǒng)優(yōu)化算法。得益于多尺度通道混洗深度可分離卷積的利用,本文算法在保證去雨質(zhì)量的同時大大降低了網(wǎng)絡(luò)參數(shù),提升了模型推理效率。

表2 不同方法參數(shù)量與測試時間對比Table2 Comparison of parameters number and average running time

為了直觀地說明本文算法的改進(jìn)效果,圖6給出了2幅合成圖像的去雨結(jié)果。傳統(tǒng)優(yōu)化算法DSC和GMM很難去除雨紋,背景仍然難以區(qū)分。RESCAN和SPANet中產(chǎn)生了較大偽影,ReHEN和UMRL丟失了邊緣紋理,細(xì)節(jié)保持較差。與之對比,本文算法去掉了最多的雨條紋,同時保留了豐富的場景細(xì)節(jié),在視覺上更令人滿意。

圖6 不同算法合成數(shù)據(jù)集去雨結(jié)果可視化對比Fig.6 Visualized comparison of different rain removal agorithms on synthetic dataset

為了驗證本文算法的泛化能力,圖7展示了不同去雨算法處理真實降雨圖像的結(jié)果。圖像1中雨條紋較小,基于傳統(tǒng)算法中的GMM結(jié)果使圖像產(chǎn)生大范圍的模糊,邊緣細(xì)節(jié)保留較差,DSC結(jié)果中有殘留雨條紋,基于深度學(xué)習(xí)的算法取得了相近的去雨效果。圖像2中雨條紋較大,與其他算法相比,本文算法在去除了幾乎所有雨條紋的同時保留了較好的邊緣細(xì)節(jié)。

圖7 不同算法真實數(shù)據(jù)集去雨結(jié)果可視化對比Fig.7 Visualized comparison of different rain removal algorithms on real dataset

為了探究階段數(shù)量對網(wǎng)絡(luò)去雨效果的影響,在Rain100H數(shù)據(jù)集上采用相同策略訓(xùn)練了包含1～7個階段的網(wǎng)絡(luò)。表3列出了不同網(wǎng)絡(luò)的去雨效果?？梢钥闯?隨著去雨階段從1增加到6,去雨效果改善明顯,從6個階段增加到7個階段,只有SSIM增加0.1%,達(dá)到了網(wǎng)絡(luò)的瓶頸,為了降低過擬合風(fēng)險,兼顧去雨質(zhì)量和運行效率,設(shè)定默認(rèn)包含6個階段。

表3 階段數(shù)量對去雨效果的影響Table3 Effect of stage number on rain removal

4 結(jié) 論

本文提出一種基于輕量化設(shè)計的稠密連接網(wǎng)絡(luò)針對單幅圖像進(jìn)行去雨,從雨條紋檢測、雨條紋去除及網(wǎng)絡(luò)輕量化3方面依次展開,具體貢獻(xiàn)如下:

1)利用基于多尺度聚集自適應(yīng)雨條紋檢測模塊對雨條紋進(jìn)行精準(zhǔn)定位,生成的注意圖用以引導(dǎo)后續(xù)模塊對雨條紋區(qū)域給予有意識的關(guān)注;

2)針對圖像去雨任務(wù)將稠密連接網(wǎng)絡(luò)進(jìn)行改進(jìn),通過去除歸一化層增強網(wǎng)絡(luò)泛化能力、使用空洞卷積增加網(wǎng)絡(luò)感受野、利用擠壓-激勵網(wǎng)絡(luò)增強表征能力,然后使用循環(huán)策略實現(xiàn)雨條紋的去除;

3)提出多尺度通道混洗的深度可分離卷積,替代稠密連接網(wǎng)絡(luò)中的常規(guī)卷積,在保證去雨性能的前提下大幅度降低了網(wǎng)絡(luò)參數(shù)。與現(xiàn)有算法相比,本文算法在合成數(shù)據(jù)集和真實數(shù)據(jù)集上都取得了較好的去雨效果,并且運行效率高,驗證了算法的效性。