基于聯(lián)合注意力機制的單幅圖像去雨算法

徐成霞，閻慶，李騰，2，苗開超

（1.安徽大學電氣工程與自動化學院，合肥 230601；2.計算智能與信號處理教育部重點實驗室（安徽大學），合肥 230601；3.安徽省氣象局安徽省公共氣象服務(wù)中心，合肥 230091）

0 引言

近年來，深度學習在計算機視覺中被廣泛應(yīng)用，在圖像識別［1］、圖像去噪［2］等領(lǐng)域取得巨大成功。Fu 等［3］結(jié)合深度殘差網(wǎng)絡(luò)［4］，建立無雨圖像和雨圖的負向殘差映射，縮小映射范圍，優(yōu)化映射空間，并通過實驗驗證了算法的去雨能力。Zhang 等［5］搭建一種自動確定雨線密度信息的多流密集網(wǎng)絡(luò)，根據(jù)雨線密度評估出相應(yīng)的標簽，根據(jù)標簽信息指導去雨任務(wù)。算法在有標簽的場景中取得良好的效果，但對于現(xiàn)實給出的標簽存在偏差。Wang 等［6］構(gòu)建含有29.5×103對真實降雨數(shù)據(jù)集，提出基于空間注意力的圖像去雨網(wǎng)絡(luò)（Spatial Attentive Network，SPANet），算法利用空間注意力從局部到全局定位雨線信息，重構(gòu)無雨背景。然而單一的空間注意力忽略了特征通道的信息不一致性，無法有效抑制冗余信息，突出關(guān)鍵特征。

基于此，本文提出通道注意力和空間注意力的級聯(lián)模式，利用通道注意力對不同通道賦予不同權(quán)重，實現(xiàn)對特征通道的抑制或突出，同時聯(lián)合空間注意力對特征通道簇進行差異性的特征提取。其次，使用改進的深度殘差收縮模塊替換原有的標準殘差模塊，通過軟閾值子網(wǎng)絡(luò)的自動學習以逐漸抑制數(shù)據(jù)集中的噪聲和冗余信息，提升算法對不同降雨場景的適用性。實驗結(jié)果表明，相較于主流去雨算法，本文算法的去雨效果更佳。此外，本文構(gòu)建了多場景的降雨數(shù)據(jù)集（Multiple Scenes Rain，MSR），補充降雨場景的多樣性。

1 構(gòu)建MSR

數(shù)據(jù)集是深度學習取得成功的重要因素，目前去雨算法大多為有監(jiān)督的訓練方式，輸入與監(jiān)督信號為有雨/無雨圖像對。Yang 等［7］構(gòu)建的Rain100L&R100H，訓練集有2 000對，測試集有200 對。Zhang 等［8］構(gòu)建的Rain800 中訓練集700 對，測試集100 對。Zhang 等［5］提出更大規(guī)模的合成數(shù)據(jù)集DID-MDN（Density-aware single Image De-raining using a Multi-stream Dense Network），數(shù)據(jù)集包含3 種密度的雨線，共12 000 對。合成場景添加的噪聲存在規(guī)律性較強、雨線分布單一、噪聲痕跡明顯等問題，去雨算法在訓練時容易對相同噪聲產(chǎn)生依賴性，不適用于規(guī)律較為復(fù)雜的真實場景。

Wang 等［6］提出半自動的算法，結(jié)合雨線的時間特性和人為監(jiān)督，在靜態(tài)背景的降雨序列中獲取無雨圖像，為去雨任務(wù)提供了很大的幫助；但其場景不豐富，圖像的視野較小，對于數(shù)據(jù)驅(qū)動的深度學習算法，不能有效地訓練去雨算法歸納大場景的降雨規(guī)律，導致雨線殘留。針對上述問題，本文構(gòu)建了MSR，如圖1 所示。

圖1 MSR雨圖示例Fig.1 Examples of rain image in MSR

MSR 合成數(shù)據(jù)集分為三步：首先生成不同密度的隨機噪聲以模擬不同大小的雨量；其次將噪聲拉長、旋轉(zhuǎn)方向，模擬雨線的多樣性；最后將生成的雨滴噪聲和原始圖像疊加，合成雨圖。雨線密度為輕度、中度、重度；雨線傾斜角度方面，以垂直方向為0°，分別有向右傾斜5°、10°、15°、20°、25°、30°，向左傾斜5°、10°、15°、20°、25°、30°，以及部分角度的混合疊加，共計20 種不同方向的合成雨線。合成降雨數(shù)據(jù)集共有6 456 張無雨圖片，32 280 張合成雨圖，圖片的分辨率為256×256。MSR 多樣性的合成雨線，為算法提升去雨精度提供數(shù)據(jù)保障。此外，MSR 通過引用文獻［6］中的半自動算法，在大量靜態(tài)背景降雨序列中獲取真實降雨數(shù)據(jù)集，共11 437對，圖像分辨率為256×256。不同的是，MSR 真實降雨數(shù)據(jù)集大多采用視野寬闊的場景，并且不對數(shù)據(jù)集進行過度切割，避免造成場景視野較小，無法有效幫助去雨算法在較大場景中的精確去雨的問題。

2 聯(lián)合注意力去雨算法

2.1 深度殘差收縮網(wǎng)絡(luò)

圖像去雨任務(wù)中，雨圖經(jīng)常含有雨線之外的噪聲，例如光暈和霧氣，會模糊圖像細節(jié)，破壞場景的背景重構(gòu)，造成較為嚴重的濾波效應(yīng)，因此有效地抑制此類噪聲成為提升去雨算法性能的途徑之一。

軟閾值為信號降噪算法常用技術(shù)，它將小于設(shè)定閾值的信號置0。而軟閾值的設(shè)置需要通過不斷調(diào)試，缺乏擴展性。Zhao 等［9］提出將軟閾值學習與殘差模塊結(jié)合，形成殘差收縮模塊，如圖2（a）所示。網(wǎng)絡(luò)在高噪聲的振動信號特征學習中，自適應(yīng)調(diào)整閾值，有效地抑制噪聲信息，提升網(wǎng)絡(luò)的特征學習能力。

圖2（b）為軟閾值學習模塊，首先對殘差單元最后一層的輸出進行絕對值化，經(jīng)全局平均池化（Global Average Pooling，GAP）獲取與上一層相同卷積核數(shù)量的一維向量，通過2 層全連接網(wǎng)絡(luò)和激活函數(shù)，將其參數(shù)縮放至（0，1）內(nèi)，其公式為：

圖2 殘差模塊結(jié)構(gòu)Fig.2 Residual module structure

其中：zc為第二層全連接網(wǎng)絡(luò)的第c個神經(jīng)元。輸出αc為相應(yīng)的縮放參數(shù)，其閾值如下：

其中：τc表示特征圖第c個通道的閾值；w、h分別為特征圖的寬和高。將縮放參數(shù)α與|x|的均值相乘，得到閾值。其目的是將軟閾值控制在合理范圍內(nèi)，避免輸出特征全為零。原深度殘差收縮網(wǎng)絡(luò)適用于一維信號，本文對其進行改進，使其應(yīng)用于圖像數(shù)據(jù)。后續(xù)實驗結(jié)果表明，深度殘差收縮模塊的加入，提升了網(wǎng)絡(luò)對于光暈、霧氣等噪聲抑制作用，有效地緩解此類噪聲對于背景重構(gòu)的影響。

2.2 聯(lián)合注意力去雨網(wǎng)絡(luò)

2.2.1 聯(lián)合注意力機制

去雨算法需要檢測雨線位置信息，定位雨水需要擴大感受野，獲取上下文信息。目前用于定位雨水的模塊主要有擴張網(wǎng)絡(luò)［10］、非局部模塊［11］和注意力機制［12］，本文采用聯(lián)合注意力機制定位雨線。注意力機制可以提升網(wǎng)絡(luò)篩選信息的效率，抑制不相關(guān)信息。對于主干網(wǎng)絡(luò)的特征圖F∈RC×H×W，其中，C、H、W依次為特征圖通道數(shù)、高度、寬度。如式（3）和式（4）所示，注意力機制將特征F通過通道注意力機制Mc∈RC×1×1得到F′，將F′通過空間注意力機制Ms∈R1×H×W得到F"，?為矩陣逐元素相乘，如圖3（a）所示。

通道注意力將特征圖在空間維度上進行壓縮，獲取一維矢量，如圖3（b）所示。其中最大池化（Max Pooling，MaxPool）和平均池化（Average Pooling，AvgPool）用于聚合特征映射的空間信息，將信息輸入到共享網(wǎng)絡(luò)，壓縮輸入特征的空間維度，得到，逐元素求和，產(chǎn)生通道注意力圖：

其中：σ為激活函數(shù)；MLP為共享網(wǎng)絡(luò)；W0、W1分別為共享網(wǎng)絡(luò)的第一層和第二層。

空間注意力如圖3（c），將通道注意力模塊的特征圖作為輸入，在通道維度分別進行了平均值池化和最大值池化，得到，將它們進行通道疊加，再依次經(jīng)過卷積、激活函數(shù)，產(chǎn)生空間注意力圖：

圖3 注意力模塊Fig.3 Attention mechanism modules

其中：f7×7代表卷積核尺寸7×7 的卷積層。對于中間特征圖，會沿著兩個獨立的維度依次推斷注意力圖，然后將注意力圖與輸入的特征圖相乘以進行自適應(yīng)特征優(yōu)化。

2.2.2 聯(lián)合注意力去雨網(wǎng)絡(luò)

聯(lián)合注意力模型（Joint Attentive Model，JAM）利用通道注意力對不同特征通道中雨線信息的檢測，賦予特征通道不同的權(quán)重，表征某通道與關(guān)鍵信息的關(guān)聯(lián)度。收集各通道與雨線的關(guān)聯(lián)度后，使用空間注意力以局部至全局的方式檢測每個特征通道的雨線及鄰域的位置信息，進行雨線的詳細定位，如圖4（d）。JAM 在主干網(wǎng)絡(luò)中采用四向循環(huán)神經(jīng)網(wǎng)絡(luò)（Recurrent Neural Networks with ReLU and Identity matrix initialization，IRNN）層［14］，IRNN 通過疊加2 個循環(huán)神經(jīng)網(wǎng)絡(luò)（Recurrent Neural Network，RNN）層［15］積累大范圍多變的信息。通過信息在整張圖中進行有效的傳播，改善無規(guī)律的真實雨線的檢測效果。

圖4 聯(lián)合注意力去雨網(wǎng)絡(luò)架構(gòu)Fig.4 Joint attention de-raining network architecture

本文在深度殘差收縮模塊的基礎(chǔ)結(jié)構(gòu)上嵌入聯(lián)合注意力結(jié)構(gòu)，構(gòu)成聯(lián)合注意力殘差收縮模塊（Joint Attentive Residual Shrinkage Block，JARSB）。在抑制光暈和霧氣等噪聲的前提下，利用聯(lián)合注意力圖聚焦重要特征，提升網(wǎng)絡(luò)對于信息的篩選性能。聯(lián)合注意力單元（Joint Attentive Unit，JAU）由JAM 和JARSB 組成，其中JAM 生成的注意力圖，引導JARSB 通過負向映射消除雨線。在聯(lián)合注意力網(wǎng)絡(luò)（JOint Attention Network，JOANet）中，網(wǎng)絡(luò)重復(fù)使用4 次JAU 結(jié)構(gòu)，來自于JAM 的注意力權(quán)重共享，最后通過兩個深度殘差收縮模塊和一層卷積重構(gòu)去雨后的背景。

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

3.1 實驗環(huán)境和評價指標

本文在公開數(shù)據(jù)集Rain100H&Rain100L［7］、DID-MDN［5］、SPANet-Data［6］和本文合成的MSR 上進行實驗，并與目前主流單幅圖像去雨算法進行比較。

實驗采用結(jié)構(gòu)相似性（Structural SIMilarity，SSIM）［16］、峰值信噪比（Peak Signal-to-Noise Ratio，PSNR）［17］、基于像素的視覺信息保真度（Pixel-based Visual Information Fidelity，VIFP）［18］作為評價指標。算法使用PyTorch 實現(xiàn)，并在4 塊NVIDIA 1080Ti 顯卡上運行。算法訓練采用Adam 優(yōu)化器，單個訓練批次的樣本數(shù)量根據(jù)GPU 的顯存上限設(shè)為16。初始學習率設(shè)為10-5，訓練終止條件為組合損失停止下降。

3.2 實驗結(jié)果分析

3.2.1 注意力圖的對比與分析

實驗將驗證注意力機制改進的效果，實驗采用SPANet-Data 為訓練數(shù)據(jù)集，對由SPANet 和本文網(wǎng)絡(luò)生成的注意力圖進行對比。圖5（a）表示某位置為噪點的概率圖，深藍色到深紅色代表該點像素為雨滴的可能性由低變高。

由圖5（c）、（d）中紅色方框區(qū)域可知：圖5（d）的紅色方框檢測的雨線密度大于圖5（c）中紅色方框的雨線密度，佐證聯(lián)合注意力在檢測雨線方面優(yōu)于SPANet 的空間注意力。表明通過對每個特征通道包含雨線信息的分析，進行自適應(yīng)的雨線檢測，可以更有針對性、全面地檢測雨線的位置信息，指導網(wǎng)絡(luò)完成去雨任務(wù)。

圖5 空間注意力與聯(lián)合注意力的雨線檢測結(jié)果Fig.5 Rain streak detection results for spatial attention and joint attention

3.2.2 去雨算法對比實驗及分析

將本文算法與其他主流去雨算法進行對比，合成數(shù)據(jù)集隨機選用了Rain100H 和DID-MDN 中50%的數(shù)據(jù)集，真實降雨數(shù)據(jù)集采用SPANet-Data；并且為了驗證實驗中去雨算法在各種場景中去雨的魯棒性，測試數(shù)據(jù)集由合成降雨數(shù)據(jù)集Rain100H、DID-MDN 和真實降雨數(shù)據(jù)集SPANet-Data 組成，表1 中給出了各算法詳細的去雨結(jié)果的評價指標。如表1 所示，基于聯(lián)合注意力機制的單幅圖像去雨算法在DID-MDN和SPANet-Data 上均有優(yōu)異的表現(xiàn)，尤其在SPANet-Data 中，3 種評價指標均是最高；但在Rain100H 的效果一般，原因在于本文網(wǎng)絡(luò)的優(yōu)勢在于使用聯(lián)合注意力機制檢測雨線，而Rain100H 中，合成雨線的痕跡較重，嚴重破壞了背景中的細節(jié)特征，導致重構(gòu)圖像的質(zhì)量一般。

表1 不同去雨算法結(jié)果對比Tab.1 Comparison of results of different de-raining algorithms

不同去雨算法的可視化結(jié)果如圖6 所示。

圖6 不同去雨算法的可視化結(jié)果Fig.6 Visualization results of different de-raining algorithms

3.2.3 MSR驗證實驗

對于數(shù)據(jù)驅(qū)動的去雨算法，訓練數(shù)據(jù)集的質(zhì)量對結(jié)果會產(chǎn)生不可忽略的影響。此部分為驗證數(shù)據(jù)集對去雨算法效果的影響，采用MSR-Synthetic、Rain100H&Rain100L、DIDMDN 訓練本文提出的去雨算法，測試結(jié)果如表2 所示。

表2 合成降雨數(shù)據(jù)集上的交叉驗證結(jié)果Tab.2 Cross-validation results on synthetic rain data

從表2 中發(fā)現(xiàn)本文數(shù)據(jù)集MSR-Synthetic 在Rain100H 上的客觀評價指標低于Rain100H&Rain100L 的去雨算法。原因有二：

其一，MSR-Synthetic 并沒有Rain100H 同類型的噪聲，雖然雨線的檢測和去除部分良好，但是存在去雨重構(gòu)圖像過于平滑問題；

其二，Rain100H 中人工添加的雨線噪聲過于密集，嚴重影響了圖片背景信息的表達，細節(jié)信息明顯不足，對比圖7的Rain100L 的測試數(shù)據(jù)集可驗證，同樣的去雨算法和背景，對于噪聲相對較輕的雨圖，去雨重構(gòu)后保留的細節(jié)也相對較多。

圖7（b）、（c）、（d）分別為JOANet 使用Rain100H&100L、DID-MDN、MSR-Synthetic 作為訓練數(shù)據(jù)集獲取的去雨算法。交叉驗證中，MSR-Synthetic 在Rain100L、DID-MDN、MSRSynthetic 合成數(shù)據(jù)集中測試評價指標相對較高。因為這些降雨數(shù)據(jù)集的噪聲分布較為均勻，噪聲程度較輕，在原始雨圖上保留了較多的背景信息。其中，在DID-MDN 上的測試結(jié)果優(yōu)于DID-MDN 作為訓練集得到的去雨算法的測試結(jié)果。因為MSR-Synthetic 的噪聲結(jié)構(gòu)與DID-MDN 類似，且數(shù)據(jù)集種類與數(shù)量要優(yōu)于后者，所以MSR-Synthetic 在DIDMDN 上的測試評價指標高于DID-MDN 的去雨算法在本身數(shù)據(jù)集上的測試指標。

圖7 合成降雨數(shù)據(jù)集上的去雨結(jié)果Fig.7 De-raining results on synthetic rain data

本文對真實降雨數(shù)據(jù)集也做了交叉驗證，結(jié)果如表3 所示。MSR-Real 的去雨算法在SPANet-Data 和MSR-Real 中的測試指標均有優(yōu)異的表現(xiàn)，尤其在SPANet-Data 數(shù)據(jù)量遠大于MSR-Real 的情況下，在MSR-Real 中的測試指標依舊高于SPANet-Data 去雨算法的測試指標。原因主要為SPANet-Data 與MSR 真實降雨數(shù)據(jù)集的風格不一，前者將一幀圖片分割為若干子圖，若輸入尺寸維持在256×256 時，雨滴相對變大，而MSR 采用的訓練數(shù)據(jù)集為完整的一幀圖片，有利于訓練網(wǎng)絡(luò)在真實環(huán)境檢測雨線的能力?？梢暬Ч妶D8。

圖8 真實降雨數(shù)據(jù)集上的去雨結(jié)果Fig.8 De-raining results on real rain data

表3 真實降雨數(shù)據(jù)集上的交叉驗證結(jié)果Tab.3 Cross-validation results on real rain data

3.2.4 殘差收縮模塊降噪實驗

針對部分去雨任務(wù)中光暈等噪聲破壞場景的背景重構(gòu)，導致濾波效應(yīng)的問題，本文去雨算法采用殘差收縮模塊替換原有的標準殘差塊。為了客觀評價兩種殘差塊影響去雨算法的程度，實驗中基礎(chǔ)架構(gòu)采用SPANet，將SPANet 中的標準殘差塊替換為殘差收縮模塊。訓練數(shù)據(jù)集采用SPANet-Data，測試數(shù)據(jù)集來源于網(wǎng)絡(luò)數(shù)據(jù)。

對比圖9（a）、（b）分析發(fā)現(xiàn)：圖（b）重構(gòu)圖像存在局部細節(jié)模糊，主要原因為該場景中存在強光源，造成雨線周圍的細節(jié)信息被干擾，不易提取，從而對重構(gòu)背景的效果造成影響。圖9（c）圖像在保持原有圖像細節(jié)盡量不丟失的條件下，實現(xiàn)圖像去噪。由此驗證殘差收縮模塊有助于提升算法在學習過程中對于圖像噪聲的抑制作用。

圖9 殘差收縮模塊降噪實驗Fig.9 Denoising experiment of residual shrinkage module

3.2.5 消融實驗

本文分析了網(wǎng)絡(luò)各模塊對于去雨算法的影響，如表4 所示。其中：Ba代表空間注意力模塊（Spatial Attentive Module，SAM）與標準殘差塊（Residual Module，RB）的組合，Bb代表SAM 與殘差收縮模塊（Residual Shrinkage Module，RSB）的組合，Bc代表通道注意力模塊（Channel Attentive Module，CAM）與RSB 的組合，Bd代表SAM、CAM 與RSB 的組合。

表4 網(wǎng)絡(luò)模塊分析Tab.4 Network module analysis

訓練和測試的數(shù)據(jù)集均為MSR 真實降雨數(shù)據(jù)集，測試樣本共1 000 對圖像對。結(jié)合表4，由Ba和Bb可知，深度殘差收縮塊相較于標準殘差塊，其PSNR 有小幅提升，說明深度殘差收縮塊在標準殘差塊進行深度特征提取的基礎(chǔ)上，進行了噪聲抑制，消除冗余信息，一定程度上加強了背景重構(gòu)部分的細節(jié)特征。結(jié)合Bd分支代表的網(wǎng)絡(luò)結(jié)構(gòu)在對比實驗中的表現(xiàn)，佐證本文提出的基于聯(lián)合注意力的去雨算法，不僅可以利用聯(lián)合注意力精準地定位雨線信息，而且可以改善雨圖重構(gòu)的細節(jié)特征。Bd相較于Ba在測試數(shù)據(jù)集上的表現(xiàn)，PSNR 提升4.5%，SSIM 提升0.3%，驗證了本文算法提高圖像去雨的可行性。

此外，本文在消融實驗環(huán)節(jié)添加運行時間復(fù)雜度分析，測試環(huán)境不變。結(jié)合表4，針對不同分辨率在不同算法架構(gòu)中的運行時間可知，基礎(chǔ)框架Ba在兩種分辨率雨圖中的運行時間分別為0.103 s 和0.341 s，將標準殘差塊替換后，運行時間分別上升了22.33%、24.05%，其本質(zhì)是因為深度殘差收縮塊在標準殘差塊的基礎(chǔ)上增加軟閾值學習模塊抑制噪聲信息，由此增加網(wǎng)絡(luò)的學習成本。對比Bb、Bc可知，添加通道注意力機制的學習成本低于空間注意力機制的學習成本，其原因是兩個注意力機制屬于不同維度：通道注意力屬于通道維度，它更關(guān)注特征通道間的相關(guān)性，基于特定的任務(wù)，學習不同通道的重要性；空間注意力屬于尺度空間維度，關(guān)注的是單個特征通道內(nèi)不同區(qū)域?qū)τ谌蝿?wù)的貢獻度，其粒度比通道注意力細致，故學習成本較高。從整體觀測，JOANet 在不同分辨率中的運行時間相比基礎(chǔ)框架Ba分別增加了28.16%、29.33%，相較于算法在去雨精度上的提升，在時間學習成本上還需要進一步優(yōu)化。

4 結(jié)語

本文針對單幅圖像去雨，提出了基于聯(lián)合注意力的卷積神經(jīng)網(wǎng)絡(luò)，關(guān)聯(lián)不同維度的注意力機制，提出利用通道注意力捕捉各通道雨線的密集程度；同時，利用空間注意力以局部到全局的方式積累上下文信息，實現(xiàn)精準去雨。此外，本文引入深度殘差收縮模塊并對其改進，降低部分圖片中非雨線類型噪聲對于背景重構(gòu)的影響。根據(jù)公開數(shù)據(jù)集以及本文構(gòu)建的數(shù)據(jù)集中的實驗結(jié)果顯示，本文算法的去雨性能優(yōu)于其他算法，但本文去雨算法存在一定局限性，學習時間成本需要進一步優(yōu)化，這也是今后需要改善的方向。