基于高效通道注意力的多階段圖像去雨網絡

李國金, 張書銘, 林 森, 陶志勇

(1.遼寧工程技術大學電子與信息工程學院,遼寧 葫蘆島 125000;2.沈陽理工大學自動化與電氣工程學院,沈陽 110000)

0 引言

在雨天,雨滴和雨紋對背景信息的覆蓋會導致圖像變得模糊,這將嚴重降低計算機視覺任務的性能,例如圖像分割[1]、語義分割[2]和目標檢測[3]等。因此,近年來圖像去雨算法備受國內外關注,如何有效地去除雨圖中的雨紋,同時保留圖像背景信息,是一個具有挑戰性的任務。

現有的圖像去雨算法可分為兩類:一是基于模型驅動的算法;二是基于深度學習的算法。早期雨紋去除主要是采用基于模型驅動的算法,通過構建去雨模型和對圖像雨紋施加先驗約束來解決圖像去雨問題。文獻[4]通過對背景層和雨層都施加基于高斯混合模型的先驗約束,可以適應各種背景以及雨紋的外觀;文獻[5]使用自動圖像分解框架將圖像分解為高頻和低頻兩部分,然后使用字典學習從高頻部分去除雨紋;文獻[6]考慮到雨紋的線狀特性和角度分布,在基于先驗約束的去雨網絡模型中引入任意方向的梯度算子,提高了網絡的去雨泛化能力。然而,基于物理模型的去雨算法只能提取雨紋層和背景層的淺層信息,不能滿足計算機視覺任務的要求。

近年來,隨著深度學習的快速發展及其在計算機視覺任務中取得的成功,研究人員提出許多基于深度學習的圖像去雨算法。文獻[7]考慮到普通卷積無法捕獲大范圍的情境信息,將空洞卷積引入圖像去雨網絡,擴大感受野,并使用雨圖的上下文信息去除雨紋;文獻[8]提出一種多尺度去雨網絡,利用多尺度和多分辨之間的互補信息表示雨紋,進而完成去雨;文獻[9]通過一個端到端的門控上下文聚合網絡融合來自不同層次的特征,使用平滑擴張卷積消除擴張卷積引起的網格偽影;文獻[10]聯合圖像去雨中的雨紋檢測和去除,設計了一種多任務聯合雨水檢測和去除的圖像去雨網絡,使用二值映射標出雨水的位置,先檢測雨水位置再估計雨紋,最后提取背景層,去除雨紋;文獻[11]提出一種半監督學習框架,在該框架中同時從已標記和未標記的數據中進行學習,以達到圖像去雨的目的。盡管與傳統的模型驅動算法相比,基于深度學習的圖像去雨算法已經取得了更好的效果,但是仍然不能在去除雨紋的同時更好地保留圖像背景細節,導致輸出的圖像無法達到較好的視覺效果。

為了解決上述問題,本文提出一種基于高效通道注意力的多階段圖像去雨網絡。

1 本文算法

1.1 網絡模型

本文提出了一個端到端多階段去雨網絡,主要由3個部分構成,如圖1所示。首先,網絡使用3×3卷積(通道數C為40)提取雨圖的淺層特征并傳遞給高效通道注意力模塊,為不同的特征通道分配不同的權重,使網絡專注于具有重要信息的通道;然后,傳遞給3個并行階段,對于第1階段和第2階段,先通過U-Net結構的編碼-解碼器進行多尺度操作,由于不同尺度下雨紋的精細程度不同,在多尺度操作下可以更加精確地去除不同粗細的雨紋;隨后,使用Transformer 模塊篩選出有用的特征信息傳遞到下一階段;最后,在第3階段使用初始分辨率網絡(Initial Resolution Net,IRN)代替編碼-解碼器,在沒有任何下采樣操作的情況下,只對輸入圖像的原始分辨率進行處理,以確保在最終輸出的圖像中保留所需的精細紋理。此外,在不同階段之間還加入了特征融合網絡(Feature Fusion Net,FFN),該模塊利用前一階段的多尺度特征幫助補充下一階段的多尺度特征。

圖1 本文網絡模型

1.2 高效通道注意力模塊(HECAB)

在圖像去雨任務中,為不同的特征通道分配不同的權重,對于捕捉和去除雨紋非常重要。文獻[12]提到的SE(Squeeze and Excitation)通道注意力模塊通過全連接層學習特征通道權重,將減少特征圖的通道維度,對通道注意力的預測產生負面效果。因此,本文提出一個高效通道注意力模塊,如圖2所示。

圖2 高效通道注意力模塊

高效通道注意力模塊使用一維卷積代替全連接層,有效避免了因通道維度縮減造成的信息丟失,能夠捕捉到更多的雨紋信息,保護了圖像的背景細節,并且一維卷積的局部跨通道交互比全連接運算的全通道交互具有更少的參數量,提高了運算速度。

實現過程為:首先,輸入特征x1經過由卷積和批量歸一化(Batch Normalization,BN)組成的并行分支,用來提高模型表征能力。然后,使用全局平均池化(Global Average Pooling,GAP)對特征圖F1進行空間特征壓縮,特征圖從[HWC]的矩陣變成[1 1C]的特征向量,其中,H為特征圖的高,W為特征圖的寬,通道數C為80。最后,特征向量通過一維卷積處理,避免了因全連接運算導致的通道維度降低,有效保護了通道信息的傳遞。

并行分支的輸出為

F1=β(Conv3×3(x1))+β(Conv1×1(x1))

(1)

式中:x1為并行分支的輸入;β為批量歸一化操作;Conv3×3和Conv1×1分別為卷積核為3×3和1×1的卷積運算。

高效通道注意力模塊的最終輸出為

FH=δ(σ(Conv1D(α(F1)))⊙F1)

(2)

式中:δ為ReLU激活函數;σ為Sigmoid激活函數;Conv1D為一維卷積運算;α為全局平均池化;⊙為元素乘法。

在HECAB中,卷積核大小k決定了局部交互的范圍,與通道數C呈非線性關系,由于人工微調容易浪費計算資源,因此使用自適應函數[13]來自動調整卷積核k的大小,算式為

(3)

式中:|t|odd表示距離t最近的奇數;系數γ和b在本文中設置為2和1,用于改變通道數C和卷積核大小k之間的比例。

1.3 編碼-解碼器

在單一尺度下提取粗細不同的雨紋信息較為困難,為了捕捉更加全面的雨紋特征,在網絡中加入了U-Net結構的編碼-解碼器(Encoder-Decoder),如圖3(a)所示。編碼-解碼器分為下采樣和上采樣兩個部分。首先,編碼器通過下采樣逐漸將高分辨輸入映射到低分辨率,每次下采樣后特征圖尺寸變為采樣前的一半;然后,解碼器使用雙線性插值法進行上采樣處理,有助于減少輸出圖像因使用轉置卷積上采樣而出現的網格偽影。此外,在跳躍連接處加入了HECAB,為不同分辨率下特征圖的通道分配權重,提取不同分辨率下的雨紋特征。

圖3 編碼-解碼器和特征融合網絡

由于編碼器重復地使用下采樣操作,難免會丟失部分特征信息。為了避免該問題,在每兩個階段之間加入了特征融合網絡,如圖3(b)所示。來自上一階段編碼-解碼器的特征被輸送到FFN中,輸入的特征首先經過兩個1×1的卷積對特征進行細化,然后把來自兩個卷積細化后的特征相加,最后傳遞到下一階段進行特征融合。

1.4 Transformer模塊

現有的多階段圖像去雨網絡大多直接把上一階段得到的去雨圖像信息傳遞到下一階段,導致向下一階段傳遞了很多無用的特征信息。受文獻[14]的啟發,本文在網絡中加入了Transformer模塊,如圖4所示。Transformer模塊由兩部分構成,即多深度卷積頭部轉置注意力(Multi-Depthwise conv head Transposed Attention,MDTA)和門控深度卷積前饋網絡(Gated-Depthwise conv Feed-forward Network,GDFN)。

圖4 Transformer模塊

在MDTA中,輸入特征首先經過層歸一化(Layer Normalization,LN),再使用1×1卷積提取通道上下文信息,3×3深度卷積提取局部空間上下文信息。卷積后的特征經過Reshape操作后得到特征矩陣Q,K和V,用于豐富局部上下文信息。然后計算Q與K的跨通道交叉協方差,即在通道維度上使用自注意力(Self Attention,SA)機制,以生成提取全局上下文信息的轉置注意力圖。最后,通過聚合局部和全局像素交互來有效處理高分辨像。MDTA的輸出FM為

FM=Conv1×1(R(V?σ(Q?K/θ)))+x2

(4)

在GDFN中,輸入特征FM先經過層歸一化處理,然后進入由1×1卷積和3×3深度卷積組成的并行分支,1×1卷積用以擴展特征通道,3×3深度卷積用以學習局部信息。并行分支的關鍵點在于下支路會通過一個非線性的GELU函數,可以將有用的特征保留,然后與上支路得到的原特征圖相乘,從而達到增強有用特征的目的。最后使用1×1卷積降低通道維度。GDFN抑制無用信息,僅允許有用的信息傳播,從而使下一階段聚焦于更重要的信息。輸出特征FG的計算過程可以表示為

FG=Conv1×1(F2⊙F3)+FM

(5)

F2=DConv3×3(Conv1×1(LN(FM)))

(6)

F3=GELU(F2)

(7)

其中:F2為上支路輸出;F3為下支路輸出;LN(·)表示層歸一化;DConv3×3表示3×3深度卷積運算;Conv1×1表示1×1的卷積運算;GELU(·)為激活函數。

1.5 初始分辨率網絡

為了保留從輸入圖像到輸出圖像的精細特征,生成信息豐富的高分辨率雨紋特征,在最后階段引入了一個初始分辨率網絡(Initial Resolution Net,IRN),如圖5所示。

圖5 初始分辨率網絡

IRN由3個初始分辨率模塊(Initial Resolution Block,IRB)組成,每個初始分辨率模塊由8個高效通道注意力模塊和1個卷積層組成,表達式為

FI=IRB(IRB(IRB(x3)))+FD

(8)

式中:IRB為初始分辨率模塊;FI和x3分別為初始分辨率網絡的輸出和輸入;FD為特征融合網絡的輸出。

1.6 損失函數

為了獲取更多的背景細節和邊緣結構,本文使用一種雙損失函數L來訓練網絡,該損失函數可以使去雨圖像生成豐富的邊緣細節,避免梯度消失。雙損失函數L由Charbonnier損失函數[15]和邊緣Edge損失函數[8]組成,算式為

L=LChar(X,Y)+λLEdge(X,Y)

(9)

(10)

(11)

其中:X和Y分別為有雨圖像和無雨圖像;LChar為Charbonnier損失函數;LEdge為Edge損失函數;Δ(·)為拉普拉斯運算;λ控制著Charbonnier損失函數和Edge損失函數在雙損失函數L中所占權重比例,本文設置為0.05,根據經驗將常數ε設置為0.003[15]。

2 實驗與分析

2.1 實驗環境和數據集

本文所提出的多階段網絡模型基于PyTorch實現,在具有24 GiB內存的NVIDIA GeForce RTX3090上進行訓練,使用SGD優化器對模型梯度進行優化,模型初始學習率為2×10-4,訓練過程中逐漸下降到1×10-6,Batch size為8,epoch為250。使用4個公開合成數據集Rain800[16],Rain12[4],Rain100L[10]和Rain100H[10]來訓練和測試網絡。Rain800包含700對訓練圖像和100對測試圖像;Rain12中包含12對有雨和無雨圖像,僅作為測試圖像;Rain100L為僅存在一種雨紋的小雨數據集,包含200對小雨訓練圖像和100對測試圖像;Rain100H為存在5種雨紋的大雨數據集,包含1800對大雨訓練圖像和100對測試圖像。

3）合理修剪。梨園要適度密植，通過合理修剪改善通風透光條件，對減輕病害發生非常重要。修剪時要剪除密擠、冗長的內膛枝，疏除外圍過密、過旺、直立生長枝條，對發病較重的樹要適當重剪。同時調整好負載，以提高樹體抗性。

2.2 實驗結果與分析

2.2.1 合成雨圖

為進一步驗證本文算法在去雨上的優勢,將本文算法與SEMI算法[11]、DiG-Com算法[6]、GCANet算法[9]、RESCAN算法[7]、JORDER算法[10]和MSPFN算法[8]的去雨性能進行比較,并使用峰值信噪比(PSNR)[17]和結構相似性(SSIM)[18]指標進行評估。對于早期的算法以及未提供訓練代碼的算法,本文直接采用提供的預訓練模型進行測試,不同算法在合成數據集上的定量對比結果如表1所示。

表1 合成數據集上的評估指標

實驗結果表明,本文算法在測試數據集Rain100H、Rain100L、Rain800和Rain12上與其他6種對比算法相比,PSNR較次優算法分別提高了0.51 dB、0.22 dB、0.36 dB、0.33 dB、SSIM分別提高了0.007、0.001、0.008、0.015。從而說明本文所提算法在去雨性能上較其他算法有更好的效果。

為了直觀地展示本文算法的性能,給出具有代表性數據集Rain100H上的實驗結果,本文算法的合成雨圖去雨效果與其他算法的視覺對比結果如圖6所示。

圖6 Rain100H數據集的測試結果

從5組圖像中可以發現,對比算法在去雨時雨紋殘留明顯,背景呈現出不同程度的虛化。進一步發現:SEMI,DiG-Com和GCANet算法無法有效去除雨紋;JORDER算法雖然去除了大部分雨紋,但背景模糊不清且有較為明顯的偽影;RESCAN和MSPFN算法雖然可以去除雨紋,但丟失了部分背景細節,并且邊緣不夠銳利;本文算法在將雨紋去除的同時,能夠較好地恢復圖像的背景細節,更加接近無雨圖像。因此,本文算法在去除雨紋和恢復背景圖像細節方面具有更好的效果。

2.2.2 真實雨圖

為進一步驗證本文算法在真實雨圖去雨中的泛化性能,將MSPFN、RESCAN、JORDER、GCANet、SEMI和本文算法的去雨表現進行對比,測試所用的模型為Rain100H上訓練的模型。由于真實雨圖沒有無雨圖像做對比,無法使用PSNR和SSIM做評估指標,本文采用自然圖像質量評估指標(NIQE)[19]作為評價指標,其值越小表明輸出圖像的質量越高,不同算法在真實雨圖上的定量結果如表2所示。

表2 真實雨圖的評估指標

表2中數據為68張真實雨圖的NIQE平均值。實驗結果表明,相比SEMI、GCANet、JORDER、RESCAN和MSPFN算法,本文算法在真實雨圖上的NIQE達到最小值,為3.677。不同算法的真實去雨視覺效果如圖7所示。

圖7 真實數據集的測試結果

由圖7可以看出,在前兩組圖像中,SEMI和GCANet算法處理后依然有大量雨紋殘留,JORDER和RESCAN算法雖然可以去除大部分雨紋,但同時也使圖像細節變得模糊,MSPFN算法仍然存在少量雨紋,部分背景細節模糊;第3組圖像中,其他對比算法均存在不同程度的雨痕和偽影,背景細節較差。通過3組真實去雨圖像對比,本文算法在去除雨紋和恢復背景細節方面效果更優。

2.3 消融實驗

2.3.1 網絡組成消融實驗

為驗證本文所提網絡結構的有效性,消融實驗針對不同網絡組成在Rain100H上的去雨結果進行分析,如表3所示。

表3 網絡組成消融實驗

本次消融實驗設計了4種不同的網絡組合方式,分別是用編碼-解碼器代替原網絡的IRNet(Net1)、用SE通道注意力模塊代替HECAB(Net2)、在原網絡中不引入Transformer模塊(Net3)和本文原網絡。

由表3分析可知:Net1中用編碼-解碼器代替網絡第3階段的初始分辨率網絡時,PSNR和SSIM分別下降了0.83 dB和0.121;Net2的PSNR和SSIM較原網絡分別下降1.12 dB和0.061,說明在圖像去雨網絡中,HECAB相比于SE通道注意力模塊具有更好的表現;Net3中不引入Transformer模塊,PSNR和SSIM分別下降了1.43 dB和0.077,這表明在多階段網絡中加入Transformer模塊抑制無用信息傳遞,網絡的去雨能力會有明顯提升。實驗結果表明,本文所使用原網絡組成在去雨上有較好的性能。

2.3.2 損失函數消融實驗

為驗證本文所使用雙損失函數L的有效性,對MSE損失函數、Edge損失函數、Charbonnier損失函數、MSE與Edge損失函數的組合、Edge與Charbonnier損失函數的組合進行實驗。損失函數的消融實驗對比結果如表4所示。

表4 損失函數消融實驗

實驗結果表明,由Edge和Charbonnier組成的雙損失函數訓練網絡的性能與單獨使用MSE、Edge和Charbonnier損失函數相比,在PSNR指標上分別提升1.27 dB、0.89 dB和1.11 dB,在SSIM上分別提升0.065、0.052和0.063。同時,比MSE和Edge組成的雙損失函數在PSNR和SSIM指標上分別提升了0.53 dB和0.027。因此,本文所使用的雙損失函數相比其他單個損失函數和雙損失函數在去雨效果上有更好的表現。

3 結論

本文提出一種基于高效通道注意力的多階段圖像去雨網絡。在網絡中使用高效通道注意力模塊為不同的特征通道分配不同的權重,使網絡更專注于具有重要信息的通道。為保留從輸入圖像到輸出圖像的精細特征,生成信息豐富的高分辨率雨紋特征,使用初始分辨率模塊代替了U-Net結構的編碼-解碼器。此外,在不同階段中還加入了特征融合模塊和Transformer模塊,阻止了無用信息進一步傳遞。與對比算法相比較,本文算法在逐漸去除圖像中雨紋的同時,也很好地恢復了背景圖像,去雨效果更優越。