用于單張圖像去雨滴的輕量級網(wǎng)絡①

蔡林紋,王 冠

(天津大學 數(shù)學學院,天津 300350)

下雨是一種常見的天氣狀況,雨不僅會影響人類的視覺感知,而且會影響室外的計算機視覺系統(tǒng).在雨天,雨線(rain streak)和附著在相機鏡頭上的雨滴(raindrop)都會顯著降低拍攝圖像的質(zhì)量.因此,雨線或雨滴去除問題引起了廣泛的研究.近些年來,研究者們在去除雨線上已經(jīng)取得了很好的成果.然而,這些方法不能直接用于去除雨滴,因為雨滴的形狀和物理成像過程與雨線完全不同.圖像上的雨線一般是呈條紋狀均勻分布在整張圖像上,且雨線的方向是一致的.與雨線相比,雨滴呈現(xiàn)出更加不規(guī)則的幾何形狀,在圖像上的分布沒有規(guī)律,使得雨滴去除問題更具有挑戰(zhàn)性.

一些研究者們基于雨滴的物理成像過程或幾何形態(tài)來模擬雨滴的分布.Kurihata 等[1]試圖通過主成分分析(PCA)來檢測雨滴,但當場景中的物體與雨滴相似時,PCA 無法有效地提取雨滴的特征,而且容易獲得錯誤的檢測結(jié)果.Yamashita 等[2]先是通過測量玻璃和立體相機之間的距離來檢測雨滴,用其它圖像區(qū)域替換該雨滴位置來達到去除雨滴的目的,但這種方法是在假設其它圖像的相同位置沒有雨滴覆蓋的前提下進行的,實際中這一假設有較大的局限性.之后,Yamashita等[3]又通過移動相機拍攝出的圖像序列來替代立體相機拍攝出的圖像.

近些年來,隨著深度學習的流行和廣泛應用,有一些研究者們開始嘗試利用神經(jīng)網(wǎng)絡來解決雨滴去除問題.Eigen 等[4]建立了一個三層網(wǎng)絡來提取合成圖像中雨滴的特征.雖然這種方法在雨滴小而稀疏的情況下效果較好,但它不能去除大而密集的雨滴.Qian 等[5]提出了一種可以同時檢測和去除雨滴的方法Attentive-GAN,該方法先利用一個小型遞歸網(wǎng)絡檢測雨滴在圖像上的位置并生成相應的注意力圖(attention maps),然后這些注意力圖進一步用于后續(xù)網(wǎng)絡處理,以提高去雨滴的性能.雖然該方法可以很好地去除單張圖像中的雨滴,然而其網(wǎng)絡結(jié)構(gòu)包含的參數(shù)相對較多,這限制了它在計算資源有限的實際應用中的潛在價值.因此,設計有效且高效的雨滴去除算法對于廣泛的實際應用是很有必要的.由于PReNet[6]在去除雨線方面是一種很好的輕量級網(wǎng)絡基本模型(Baseline),但直接將其應用于雨滴去除任務,得到的去雨滴結(jié)果并不好.

本文基于PReNet 提出了一種用于單張圖像雨滴去除的網(wǎng)絡PRSEDNet,與現(xiàn)有的基于深度學習的去雨滴方法相比,可以獲得相近或更優(yōu)的去雨滴效果.特別地,PRSEDNet 顯著降低了網(wǎng)絡的參數(shù)量.實驗結(jié)果驗證了我們的PRSEDNet 可以在達到高效的去雨滴性能的同時,有更少的參數(shù)量,能獲得和Attentive GAN 相近的去雨效果,并且比輕量級網(wǎng)絡PReNet 表現(xiàn)得更好.

1 本文方法

為了更好地細化分析本文方法,在本節(jié)中，我們首先介紹了雨滴去除問題的基本理論模型,其次對我們使用的網(wǎng)絡結(jié)構(gòu)進行了描述，最后給出了我們使用的損失函數(shù).

1.1 基本理論模型

Attentive GAN中將雨滴圖像建模為:

其中,I是帶雨滴的圖像,B和R分別為背景圖層和雨滴圖層,M是一個二值變量,當M=1時,代表此像素屬于雨滴區(qū)域,當M=0時,代表該像素屬于背景區(qū)域.但其對M的真值采用了如下的設置:當x≥30,M=1; 當x＜30,M=0,其中x為圖像像素(取值范圍為0-255).該方法簡單地將閾值設置為30,這種方式在實際應用中有一定的局限性.為了讓模型有更好的泛化能力,我們不使用雨滴位置的信息,即上述的M,而是將雨滴去除問題看作是兩個圖層的分離問題,建模如下:

其中,I,B,R的定義同上.

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

本文提出的去雨滴網(wǎng)絡PRSEDNet 如圖1所示,由兩個卷積層,卷積長短期記憶網(wǎng)絡和特征提取模塊組成.即1 個卷積層fin用于接收輸入圖像,經(jīng)過卷積長短期記憶網(wǎng)絡flstm和特征提取模塊fblock來提取特征,1 個卷積層fout用于輸出去除雨滴后的清晰圖像.因此,第t階段PRSEDNet 可以用如下公式表述:

圖1 PRSEDNet的網(wǎng)絡結(jié)構(gòu)

最大階段數(shù)為T,除特殊說明,網(wǎng)絡中所有的卷積層的卷積核大小都設置為3,步幅(stride)和填充(padding)都設置為1,激活函數(shù)使用ReLU[7].在網(wǎng)絡的整個過程中,都沒有上下采樣的操作,卷積核的其它參數(shù)設置都是為了讓特征圖的分辨率不會發(fā)生改變,旨在恢復出的清晰圖不會有圖像內(nèi)容細節(jié)的損失.下面,我們將給出每一層的實現(xiàn)細節(jié).

fin:該部分僅由一個卷積層組成.3 通道的雨滴圖y和3 通道的輸出圖xt-1按通道數(shù)連接在一起組成一個6 通道的輸入,將6 通道輸入送給卷積核大小為3的卷積層,輸出一個32 個通道的特征圖.其中,x0初 始化為y.

flstm:這部分參考PReNet的網(wǎng)絡結(jié)構(gòu),本文也使用ConvLSTM[8]中所提出的Convolutional LSTM 結(jié)構(gòu).Convolutional LSTM 接收一個32 個通道的特征圖，最終輸出一個通道數(shù)仍為32的特征圖.

fblock:此部分為特征提取模塊，一共由5 個RSED Block 組成,每個RSED Block 可以分為3 個部分,即兩個卷積層和一個注意力機制模塊,具體結(jié)構(gòu)如圖2所示.接收到的特征圖先通過一個卷積核大小為3的空洞卷積,該空洞卷積的填充(padding)和擴張(dilation)都設置為10.之后再經(jīng)過一個卷積核大小為3的卷積層(該卷積不是空洞卷積,步幅和填充均設置為1),最后再通過一個注意力機制模塊,我們在實驗中采用SENet[9]中的SE 模塊作為注意力機制模塊,SE的具體結(jié)構(gòu)如圖3所示,H和W分別代表特征圖X的長和寬,C表示通道數(shù),r是放縮因子.實驗中我們將C和r分別設置為32和16.每個RSED Block中的輸入和輸出都是32 個通道,在特征提取的整個過程中我們和文獻[10]一樣使用了殘差學習[11].

圖2 RSED block的具體結(jié)構(gòu)

圖3 SE 模塊的具體結(jié)構(gòu)

fout:該部分作為最終的輸出層僅由一個卷積層組成.該卷積層的卷積核大小為3，且沒有激活函數(shù).它將前面特征提取模塊輸出的32 個通道的特征作為輸入,輸出一個3 通道的圖像,并將其加到原始雨滴圖y上,生成最終的去雨滴圖像.

1.3 損失函數(shù)

PRSEDNet 有T個階段,通過實驗我們發(fā)現(xiàn)只需對網(wǎng)絡最終的輸出xT進行監(jiān)督.我們不再簡單地使用SSIM 損失函數(shù),而是將L1損 失和LW-SSIM損失[12]結(jié)合起來,最終采用的損失函數(shù)如下:

其中,α和β分別代表兩種損失在整個損失中的權(quán)重系數(shù),gt是用于監(jiān)督訓練的無雨滴清晰圖像,LW-SSIM是在SSIM 損失函數(shù)中運用哈爾小波[13],使SSIM 損失函數(shù)接收到的不再是整個圖像,而是經(jīng)過哈爾小波變換計算出的各個圖像切片.通過這種方式,我們可以在損失函數(shù)中改變圖像分辨率,以有效地抑制最終輸出圖像出現(xiàn)不希望的偽影.

2 實驗分析

2.1 實驗方法

實驗采用全監(jiān)督方法,用雨滴數(shù)據(jù)集(參考2.2 節(jié))中雨滴圖和相對應的無雨滴圖(ground truth)形成訓練對,然后在我們提出的網(wǎng)絡結(jié)構(gòu)上進行端到端(imageto-image)的訓練.旨在學習到雨滴圖到無雨滴圖的映射或殘差.

2.2 數(shù)據(jù)集和評價指標

實驗采用AttentiveGAN中Qian 等創(chuàng)建的雨滴數(shù)據(jù)集[5]來訓練和測試我們的算法.該數(shù)據(jù)集一共有861對雨滴圖和清晰圖用于訓練,含有兩個測試集,test A和test B.其中test A是test B的子集,test A 包含58對質(zhì)量優(yōu)良且對齊的圖像,test B 包含249對圖像,但test B中有部分圖像沒有完全對齊,用來計算評價指標會得到不準確的結(jié)果,為了便于定量比較并精準計算評價指標,我們只在test A 上測試我們的結(jié)果,并將結(jié)果與其他去雨滴算法進行比較.

評價指標我們采用最常用的衡量圖像質(zhì)量的兩個值,峰值信噪比(PSNR)[14]和結(jié)構(gòu)相似度(SSIM)[15].PSNR和SSIM 值越大,表示圖像失真越小,質(zhì)量越好.

2.3 數(shù)據(jù)預處理和實驗條件

在整個網(wǎng)絡訓練中,我們設置了合適的卷積核以及相應的步幅和填充以使處理圖像的分辨率不發(fā)生變化.為了加快網(wǎng)絡的訓練且減小顯存占用,我們將訓練的861對圖像全都裁剪為100×100 分辨率的圖像塊,然后將所有的圖像塊輸入到網(wǎng)絡進行圖像端到端的訓練.

我們設置網(wǎng)絡的最大階段數(shù)T為6.損失函數(shù)(4)中的權(quán)重系數(shù) α和β分別設置為0.15和0.85.實驗采用PyTorch 框架,并在單張NVIDIA GTX 1080Ti GPU 上進行訓練.梯度下降的優(yōu)化器采用Adam[16],Batchsize固定為18.整個實驗訓練100 個epoch,初始學習率設為0.001,當達到第30,50和80 個epoch 時,學習率變?yōu)榍耙浑A段的一半.

2.4 數(shù)值和視覺結(jié)果

為了量化去雨滴的性能,表1給出了Eigen,PReNet,AttentiveGAN,PRSEDNet 四種算法的結(jié)構(gòu)相似度與峰值信噪比的數(shù)值比較(其中上下箭頭分別代表提升和不足,數(shù)值是我們的結(jié)果和其它算法結(jié)果相減得到的).其中Eigen的結(jié)果是用其作者訓練好的模型直接在測試集test A 上進行測試,而AttentiveGAN的作者在GitHub 上并未給出訓練代碼,只給出了其訓練好的模型.我們直接使用其提供的模型在測試集test A 上進行測試.雖然PReNet的作者在GitHub 上提供了訓練好的模型,但由于它是用于處理雨線去除問題,其提供的模型不能直接用于雨滴去除問題,我們便按照其作者的默認設置參數(shù)在訓練集上重新訓練了PReNet 來進行測試.可以看到,我們的結(jié)果遠遠優(yōu)于Eigen的結(jié)果,和我們的基準模型PReNet 相比有很大的提升,同時也能達到和AttentiveGAN 相近的結(jié)果.我們還給出了PReNet,AttentiveGAN,PRSEDNet 三種算法的參數(shù)量比較,具體見表2,值得注意的是,我們比PReNet 只多了很少的參數(shù),但去雨滴效果卻提升很多,而且我們能得到和AttentiveGAN 很相近的結(jié)果,但我們的參數(shù)量只有其參數(shù)量的2.7%左右.

表1 幾種算法在測試集test A 上的PSNR與SSIM 平均值比較

表2 幾種算法的參數(shù)量比較

圖4給出了本文和其他算法在測試集test A 上的去雨滴效果圖,為了便于觀察,我們從每張圖中選取了3 個矩形區(qū)域進行放大并排列在相應圖的下方.如圖4中的第一幅圖所示,Eigen 算法只能處理很小且稀疏的雨滴,當遇到大而密集的雨滴時,其不能達到去除雨滴的效果,PReNet 能去掉部分雨滴,但其去除雨滴不徹底(參考紅色和綠色特寫框).從圖4中第2 幅圖的3 個特寫框可以看出,四種算法去除雨滴都不徹底,但Eigen,PReNet,AttentiveGAN 三種算法會在得到的清晰圖中保留更明顯的雨滴印跡,相比之下我們的算法可以得到更優(yōu)的效果.圖4中第3 幅圖的綠色特寫框反映了PReNet 在生成清晰圖像時會產(chǎn)生不希望出現(xiàn)的偽影,而我們的算法在損失函數(shù)中運用了哈爾小波變換,能有效地抑制該問題的出現(xiàn).通過視覺效果圖的對比可以看出,我們所提出算法的去雨滴效果遠遠優(yōu)于Eigen 算法的效果,同時可以達到不輸于AttentiveGAN算法的效果,并且比同為輕量級網(wǎng)絡的PReNet 效果更好.

圖4 不同算法在測試集test A 上的去雨滴效果對比

為了更好地證明我們算法的有效性,圖5給出了整個實驗流程中PRSENet和PReNet 隨epoch 變化的PSNR和SSIM 值.從圖5中可以看出，PRSENet的PSNR和SSIM 值在實驗中均大于PReNet的PSNR和SSIM 值,這進一步說明了PRSENet 在去雨滴任務中比PReNet 更有效.另一方面,PRSENet的PSNR和SSIM 曲線最終趨于平緩，這也說明了我們算法的在實驗中是可以收斂的.

圖5 訓練過程中PSNR和SSIM 值的變化

3 總結(jié)

本文提出了一種用于單張圖像去雨滴的輕量級網(wǎng)絡算法.實驗結(jié)果表明,與當前流行的去雨滴方法Eigen和AttentiveGAN 相比,我們的PRSEDNet 可以達到更優(yōu)或相近的去雨滴效果,和我們的基本模型輕量級網(wǎng)絡PReNet 相比,無論是在去雨滴清晰圖的視覺感觀上,還是在PSNR和SSIM 數(shù)值上,我們的去雨滴效果都更好.此外,我們所提出的輕量級網(wǎng)絡也可以應用于其他圖像恢復任務中,以平衡圖像的恢復質(zhì)量和相應的計算開銷.