基于改進擴散模型的圖像去雨方法

錢 楓,胡桂銘,祝 能,鄧明星,王 潔,許小偉

(武漢科技大學 汽車與交通工程學院, 武漢 430081)

0 引言

雨天環境下,車載攝像頭拍攝到的圖像有雨痕干擾,嚴重影響高級駕駛輔助系統(advanced driving assistance systems,ADAS)識別前方車輛信息的準確性,因此圖像去雨對智能駕駛具有極其重要的意義。

目前,關于雨水去除的研究主要有傳統方法和深度學習方法。傳統去雨方法主要有高斯混合模型[1]、字典學習和稀疏編碼[2-3]等,這些方法將雨圖像分為雨痕層和背景層,并通過剝離雨痕來得到背景層。然而,這些傳統去雨方法在處理復雜圖像或大雨等情況時容易存在雨痕殘留、紋理信息丟失等問題[4]。

相比之下,基于深度學習的圖像處理方法具有更好的泛化性,且可以實現端到端訓練。在卷積網絡中嵌入殘差模塊、注意力機制等模塊可使網絡更好地處理細節信息。Han等[5],Yuan等[6],柳博等[7]利用擴展的深度卷積神經網絡學習雨的殘差圖像,避免了參數估計不準確對結果的影響。隨著深度學習框架的不斷優化,Wang等[8]通過引入循環網絡并利用長短期記憶來連接各個階段,使實驗結果更優;陳崢等[9],張焱等[10]運用門控循環單元設計循環神經網絡進行去雨。此外,也有其他組合方式和網絡結構應用于圖像去雨,如生成對抗網絡(generative adversarial network,GAN)在圖像去雨領域表現出顯著優勢,朱德利等[11],Yu等[12],李然等[13]使用GAN對抗訓練來提高圖像保真度,達到優異的去雨效果。但GAN會存在訓練不穩定、模式崩潰等問題,主要原因在于生成數據和真實數據的分布重合度不高[14]。上述基于深度學習的去雨方法,使用注意力機制、改進的殘差塊等對雨痕進行去除,然而,這些方法大多未考慮雨痕類型(如雨滴、雨線)對圖像的影響,有較多的雨痕殘留,去雨質量仍有較大的提升空間。

為了解決上述圖像去雨質量差的問題,本文中提出一種基于擴散模型的圖像去雨方法,通過設計殘差模塊雙輸入信息通道和添加ECA通道注意力機制對網絡進行改進,提高了模型捕獲跨通道交互信息的能力,進而有效去除雨痕,提高了去雨質量,泛化性更好。

1 擴散模型

擴散模型作為生成模型之一,是一種依賴先驗的條件模型,能夠實現從噪聲(采樣自簡單的分布)生成目標數據樣本。模型[15]包括前向和逆向2個過程,通過在數據中加入隨機噪聲(前向過程)并從中采樣出所需的數據樣本(逆向過程)來實現,擴散過程如圖1所示。

圖1 擴散過程示意圖

1.1 前向過程

前向過程是指逐漸向數據添加高斯噪聲直至數據趨于高斯分布。在數據中包含著T個時間(T在訓練過程中為可變參數)的擴散過程,當前時間是通過上一時間的數據按照式(1)所示添加高斯噪聲。

(1)

式中:q(xt|xt-1)為先驗分布;N為高斯分布;xt為時間t對應的加噪數據;I為高斯噪聲的方差;βt在擴散過程中從0.0 001到0.02線性增大,并且加入高斯噪聲的均值和方差由βt決定。

前向過程采用預先設定好的方差調整方案,使得整個擴散是一個固定的過程。當時間T足夠大時,得到的xt則變為完全的隨機噪聲。

1.2 逆向過程

擴散模型的逆向過程為去噪過程,通過噪聲估計網絡推測真實分布q(xt|xt-1),從隨機噪聲開始逐漸去噪,最終生成一個真實的樣本。公式如式(2)所示。

p(xt-1|xt)=N[xt-1;μθ(xt,t),σ2I]

(2)

式中:p(xt-1|xt)為后驗分布;σ2為方差;μθ(xt,t)為參數化估計均值,通過噪聲估計網絡對噪聲εθ(xt,t)預測進而間接得到。

網絡訓練公式如式(3)所示,從時間t-1到0逐步重復xt還原成xt-1的步驟(即計算xt-1,將得到的xt-1作為新的xt)。

(3)

2 基于擴散模型的去雨結構

2.1 去雨模型結構

本文中運用擴散模型對圖像進行去雨處理,并改進模型的噪聲估計網絡。設計殘差模塊的雙輸入信息通道,生成更多的訓練數據,提高模型的泛化性;添加ECA通道注意力機制模塊,該模塊明確給出每個通道的權重,從而幫助理解模型的決策過程,提高模型的可解釋性。模型去雨結構如圖2所示,有雨圖至加噪數據圖(xt)過程為經過T時間加噪直至變為純噪聲的前向過程;加噪數據圖至無雨圖過程為將噪聲輸入噪聲估計網絡學習去噪從而得到無雨圖的逆向過程。

圖2 擴散模型去雨結構圖

模型的訓練過程:從有雨圖數據集內隨機選擇一個批次的樣本進行訓練,每個批次包含8個樣本,記為x0,并在1到T的時間中隨機選取一個時間t,將其對應的高斯噪聲與x0結合形成加噪后的數據結果xt。時間t能夠生成的正弦位置編碼,與xt一起送入噪聲估計網絡,準確識別當前所處在的去噪階段,預測出高斯噪聲εθ(xt,t)。通過高斯擴散計算該噪聲與隨機噪聲的損失,更新噪聲估計網絡權重,并繼續訓練直至訓練完成。

2.2 噪聲估計網絡

為了最小化實際的噪聲分布與預測的噪聲分布之間的差異,本文中采用噪聲估計網絡結構來預測噪聲并不斷更新權重。

噪聲估計網絡結構包含了下采樣模塊、中間層模塊、上采樣模塊、組歸一化(group normalization,GN)、激活函數(Swish)等。網絡模型結構如圖3所示,為了提高模型的泛化性,分別在下采樣模塊、中間層模塊、上采樣模塊中添加雙信息輸入的殘差模塊(residual module,block)中;在中間層模塊中添加ECA注意力機制。網絡模型是一種編碼-解碼(encoder-decoder)架構。編碼器的每個階段都包含下采樣模塊,用來降低特征的空間大小,然后通過解碼器將壓縮的特征逐漸恢復。為了更好地優化網絡,解碼器中使用跳躍連接,將編碼器中得到的同維度特征連接起來。

圖3 改進的噪聲估計網絡結構示意圖

網絡預測噪聲過程,將有雨圖像訓練集隨機裁剪成64×64的三通道圖片,并將其與時間t對應的正弦位置編碼一并送入網絡中進行訓練,計算t時刻后模型的噪聲圖片,可以預測出噪聲,并用預測噪聲與實際噪聲進行擬合,不斷更新權重。有雨圖像經過卷積、上采樣、中間層和下采樣等一系列特征處理,最后通過卷積操作輸出得到無雨圖。

2.2.1 雙信息輸入殘差模塊

針對網絡加深而恒等映射難以擬合的問題,引入殘差模塊作為網絡結構的基礎組成部分。本文中將殘差模塊設計為雙信息輸入,其中一個信息從訓練集中抽取雨圖作為樣本,經過組歸一化、激活函數和卷積層處理;另一個信息為從1到T時間隨機抽取的t時刻正弦位置編碼,經過激活函數和全連接層處理;而后將2個信息得到的輸出進行相加操作,從而實現網絡信息融合輸入,其網絡結構如圖4所示。該模塊中的組歸一化對各個組通道方向上的均值和方差進行歸一化,從而脫離batchsize大小的約束。在殘差模塊中加入1×1卷積(convolution,conv)調整因卷積層而改變的通道數,這種方法保留了編碼過程中不同層丟失的信息。

圖4 殘差網絡結構示意圖

模型采用殘差結構,使得特征映射對輸出的變化更為敏感,模型的梯度更容易訓練。為了保持網絡輸入輸出相同,需使h(x)=x,故令F(x)=0,公式如式(4)所示。

h(x)=F(x)+x

(4)

式中:h(x)為輸出;F(x)為殘差項;x為輸入。網絡首先從輸入中提取出特征,再與殘差結構組成新的特征。在殘差模塊中,采用跳躍映射的方式來構建殘差單元,即通過將輸入與輸出直接相加的方法,來補充在卷積過程中可能丟失的特征信息。殘差模塊可以減少數據維度,同時也降低梯度消失與梯度爆炸的風險。

2.2.2 ECA通道注意力機制模塊

注意力模塊可以定位圖像中的目標區域,有效提取圖像特征,使網絡能更好地理解雨痕區域的特征,并且在處理時保留了圖像的細節信息。為了提升模型性能,在網絡的中間層模塊加入ECA注意力機制。該機制具有較強的信息獲取能力,且不會受數據量的影響,能夠在不降維的情況下實現全局平均池化,增強整體的通道特征,提升模型性能。ECA通道注意力模塊如圖5所示,利用1×1卷積來獲取通道間的信息,避免了通道的維數下降和參數量減少的問題。

圖5 ECA模塊結構圖

ECA模塊首先對輸入的特征圖進行空間特征的壓縮,將其維度從H×W×C變為1×1×C,這一過程采用全局平均池化操作,以獲得大小為1×1×C的特征圖。然后,模塊使用1×1卷積對該特征圖進行通道特征學習,了解各通道間的重要性。整個過程,輸出的維度不變,仍然為1×1×C。最后,通過通道注意力機制,將通道注意力的特征圖1×1×C與原始輸入特征圖H×W×C逐通道相乘,從而輸出具有通道注意力的特征圖。

3 實驗結果及分析

3.1 數據集設置

本文中對2類雨天場景圖像進行實驗,一類是相機鏡頭上有雨滴干擾,另一類是圖像背景有雨線遮擋。

對于第一類的雨滴情況,采用文獻[16]提供的模擬雨滴數據集,共包含1 119對圖片,其中訓練集包含861對,測試集包含258對。對于第二類的雨線情況,使用Rain100數據集進行實驗,共包含2 000對圖片,其中訓練集包含1 800對,測試集包含200對。

3.2 訓練

3.2.1 訓練條件

本實驗的訓練使用了一臺具有24 GB內存的RTX3 090計算機,Python版本為3.9.13,使用的深度學習框架為pytorch.1.13.1。

3.2.2 訓練參數設置

本工作對GMZ07膨潤土及其不同摻砂率混合物的飽和試樣進行直剪試驗及掃描電鏡測試,研究試樣用不同濃度鹽溶液飽和時的強度特性,分析和總結試驗結果可得到以下結論.

在本文中,網絡參數設置初始學習率為 0.000 02,batchsize為8,總迭代次數設置為200 000,并且使用Adam優化器進行梯度優化。為了緩解過擬合現象,模型中引入了Dropout層,其作用是在每個訓練批次中隨機忽略50%的特征檢測器,即將一半的隱層節點值設為0,以此來降低模型的復雜度,并且可以脫離對一部分局部特征的依賴。

3.2.3 訓練過程

首先,運用本文提出的網絡分別對上述2種數據集進行迭代訓練,得到了包含雨痕信息的預訓練模型,然后,運用該模型對測試集進行測試。在訓練過程中,使用損失函數來估計雨圖中雨痕噪聲,經過約152 000次迭代,Loss曲線的斜率趨于0,并且數值接近于0。此時,可以認為模型訓練已經收斂,整個網絡達到了很好的訓練效果。

3.3 算法對比及實驗結果

本文中使用峰值信噪比(peak signal to noise ratio,PSNR)和結構相似性(structural similarity,SSIM)2項指標對去雨效果進行評價。

為了驗證本文算法的可信度,運用改進擴散模型與其他幾種去雨算法(GMM、DID-MDN、GAN、文獻[16]方法)對測試集圖片進行測試,測試得到的圖片平均PSNR和SSIM值如表1所示。使用不同方法去雨得到的部分效果圖如圖6所示。

表1 去雨方法在測試集上的平均PSNR與SSIM值

圖6 5種方法去雨結果圖

通過比較去雨操作后輸出圖像和無雨圖像之間的關系,得到PSNR和SSIM值。PSNR的計算公式如式(5)所示,其數值越高代表去雨圖的失真越小,則效果越好。

(5)

式中:MAXI為無雨圖中最大的像素值;MSE為去雨圖和無雨圖之間的均方差。

(6)

式中:x為雨圖;y為去雨圖;μx、μy分別為x、y的像素均值;σx、σy分別為x、y的像素方差;σxy為x和y的協方差;c1和c2為常數。

通過表1與圖6可知,在模擬雨滴數據集上,GMM和DID-MDN算法的去雨效果并不明顯;而GAN網絡去除了部分雨滴,但仍有雨滴殘留;另外,文獻[16]所提算法在去除雨滴效果較好,但損失了部分車輛的邊緣信息。在Rain100雨線數據集上,GMM方法去雨有部分雨痕殘留;DID-MDN僅能去除少量雨線;GAN網絡與文獻[16]提出的方法具有良好的去雨效果,但圖像的關鍵細節丟失嚴重。而改進擴散模型在去除雨滴和雨線方面均取得最高的PSNR和SSIM,其中雨滴和雨線的PSNR值分別為30.328 5和34.896 5,SSIM值分別為0.927 1和0.962 0,表現出良好的性能。通過圖6可知,本文算法能夠保留圖像的細節信息,在一定程度上減少了圖像的去雨過度,泛化性更好。

3.4 真實雨天情況下車輛檢測驗證

本實驗選擇了2個測試集與真實雨圖數據集中包含車輛的圖片,使用目標檢測算法進行車輛檢測。

為了驗證本文擴散模型去雨算法的可信度,本文中使用目標檢測算法對去雨后的真實雨滴、雨線圖像與上述雨滴、雨線測試集圖像進行車輛檢測,以目標檢測置信度作為評價指標。在當前的目標檢測器中,YOLOv7在速度和精度上都超過了以往的檢測器,故本文中使用YOLOv7進行車輛檢測。直接對4個數據集中的原始有雨圖進行車輛檢測,得到以下置信度結果:1)模擬雨滴[16]為0.49;2)Rain100為0.77;3)真實雨滴圖為0.83(白車),0.56(紅車);4)真實雨線圖為0.75(白車),0.75(黃車)。將上述5種去雨算法生成的無雨圖輸入目標檢測算法中進行檢測,不同去雨方法的檢測效果如圖7所示,計算得出的平均置信度提升值如表2所示。

表2 車輛檢測置信度提升平均值 %

圖7 5種方法去雨后檢測結果圖

由表2與圖7可知,其他4種算法在真實雨天環境下去雨后圖像背景丟失嚴重,對數據集去雨圖片的車輛檢測置信度提升有限;而本文算法在雨天情況下對車輛檢測置信度的提升均高于其他4種算法。該算法能夠通過去除雨痕來有效地提高在雨天環境下的圖像質量,從而提高后續計算機視覺任務的準確性,進一步驗證了改進擴散模型的實用性。

4 結論

針對圖像去雨任務,提出一種基于改進擴散模型圖像去雨方法。將擴散模型應用在圖像去雨領域,通過雨圖和時間來引導前向擴散加噪,從而逆向采樣得到無雨圖;并且設計殘差模塊的雙輸入數據信息和添加ECA注意力機制,提升去雨效果。

與GMM、DID-MDN、GAN和文獻[16]算法相比,在模擬雨滴和Rain100數據集去雨上,本文算法取得了最高的PSNR和SSIM值。此外,對去雨后的真實雨滴、雨線圖像與上述雨滴、雨線測試集圖像通過YOLOv7進行車輛檢測,本文所提算法對置信度的提升高于其他幾種算法,進一步驗證了本文所提算法在圖像去雨效果上優于其他方法,同時最大程度保留圖像細節部分。