基于條件擴散隱式模型單幅圖像去雨

徐 杰,孫偲遠

(江蘇科技大學 計算機學院,江蘇 鎮江 212100)

0 引 言

降雨是常見的惡劣天氣之一,從雨天獲取的圖像受到雨水的影響,會導致圖像內容模糊,從而增加了視覺任務的難度,如目標檢測[1]、語義分割[2]、場景分析[3]等等。因此,圖像去雨具有重要的研究意義。

目前有大量的去雨算法被提出。其中,傳統的圖像去雨算法主要是基于先驗信息的優化算法,如稀疏編碼[4]、高斯混合模型[5]和圖像分解[6]等等,這些算法試圖利用紋理特征的先驗信息對雨紋進行建模,然后將它們從圖像中分離出來。然而,這些方法只能提取圖像中的淺層信息,導致低級特征的模型表征能力非常有限,泛化能力和自適應能力相對較弱。因此,此類方法對于那些背景物體復雜且信息豐富的圖像會有模型建立復雜、計算量大和去雨效果不理想等問題。

隨著神經網絡的快速發展,卷積神經網絡憑借其強大的特征表示力和計算能力在圖像去雨研究上表現優異。Fu等[7]首次提出深度卷積網絡架構用于圖像去雨,這種方法將圖像分為基礎層和細節層,并通過學習細節層之間的非線性映射關系來去除雨痕,接著將細節層和基礎層合并得到去雨圖像。Zhang等[8]提出了一種基于密度感知的多流密集連接卷積神經網絡算法對雨紋進行去除,該算法將雨密度分級,并利用多流密集網絡預測出清晰圖像,具有適應多種尺度的雨紋的能力。Ren等[9]提出了一種漸進循環去雨網絡,該網絡通過多個階段的遞歸運算來實現網絡參數共享,以取代復雜模型,從而減少網絡參數并提高圖像去雨的質量。Zamir等[10]將圖像去雨分解為多個階段的子任務,其中編碼器網絡用于學習特征,最終利用恢復網絡完成圖像去雨過程。Huang等[11]提出了一種面向內存的半監督方法,使網絡能夠從數據中探索和學習雨紋的特性,并利用自監督的記憶模塊來恢復無雨背景。

在圖像去雨領域,生成網絡建模也受到廣泛關注。生成模型通過學習清除圖像背景的底層數據分布,從而實現圖像去雨。由于其更強的表達能力,生成方法在視覺恢復問題中具有更好的泛化能力。最近,去噪擴散概率模型在各種生成建模任務[12-14]中取得了顯著成功。

基于先進的生成建模方法,該文使用條件擴散概率模型,應用于處理雨天環境下的圖像恢復。為使該方法更加高效和靈活,采用基于SR3[15]的U-Net全卷積網絡架構,可支持輸入任意大小圖像;引入確定性加速采樣,用更小的采樣步數來加速生成過程,提高圖像恢復效率;通過圖像重疊分塊處理并使用平滑噪聲估計來指導采樣過程,減少運算的數據量和計算壓力,進一步提高算法的適用性。

1 基本理論

1.1 去噪擴散概率模型

去噪擴散概率模型[16]是一類通過馬爾可夫鏈學習的生成模型,能夠將高斯噪聲分布執行變分推理轉換為模型所訓練的數據分布,從而實現去噪和數據恢復的任務。擴散過程由前向擴散過程和逆擴散過程組成。在前向擴散過程中,從一個干凈的數據樣本x0開始,將預定義的方差表{β1…βt}注入高斯噪聲,βt∈(0,I),持續T個時間步長進行加噪,最終得到一系列數據樣本x1到xt。如下所示:

(1)

(2)

聯合分布pθ(x0:T)稱為逆擴散過程,具有學習高斯躍遷的馬爾可夫鏈。從標準高斯分布開始,迭代執行T個時間步來去噪,從而生成與訓練分布對應的圖像。當時間步數較大且{βt}的增量較小時,逆分布可以近似為高斯分布。利用已知逆分布q(xt∣xt-1),可以對xt～N(0,I)進行采樣,從q(x0)中得到樣本,執行逆擴散過程。使用帶參數θ的神經網絡進行近似:

(3)

pθ(xt-1∣xt)=N(xt-1;μθ(xt,t),Σθ(xt,t))

(4)

在這個模型中,使用μθ(xt,t)和Σθ(xt,t)神經網絡對逆擴散過程進行參數化。為了獲得參數θ,通過最小化定義數據分布的負對數似然的變分下界。具體而言,使用優化負數據對數似然的變分界來訓練模型。通過這種方式,可以有效地學習到數據分布的結構,提高模型的性能,并更好理解數據的特征。如下所示:

(5)

(6)

(7)

(8)

(9)

(10)

(11)

為了使分布的均值相等,訓練目標可以進一步簡化為:

Lsimple=Et,x0,[-θ(xt,t)2]

(12)

在該網絡中,首先在t時刻從xt中預測噪聲t,然后通過逆擴散過程pθ(xt-1∣xt)采樣,最終可以得到目標分布的樣本:

(13)

其中,z～N(0,I),這類似于通過朗之萬動力學采樣步驟,從噪聲擾動最大的分布中取樣,并逐步降低噪聲尺度的大小,從數據分布中推導出新的樣本,直到達到最小值。

1.2 條件擴散模型

條件擴散模型[12]展示了最先進的數據編輯和圖像合成能力。這些模型基于約束生成圖像,并且通過逆擴散過程pθ(xt-1∣xt,y)學習條件分布,而不修改前向擴散過程q(x1:T∣x0)。這種方法可以保證以y為條件的數據分布采樣的x具有高保真度。其過程可以表示為:

(14)

為了學習條件分布,在所有時間步長t均使用圖像y對神經網絡進行約束。訓練目標在此處被定義為:

Lsimple=Et,x0,[-θ(xt,y,t)2]

(15)

獲得了條件分布之后,該模型就可以用于推理。給定一幅圖像y后,可以從純高斯噪聲開始,從T步的馬爾可夫鏈中采樣,pθ(xt-1∣xt,y)的條件分布的對應均值為:

(16)

1.3 條件擴散隱式模型

去噪擴散隱式模型[16]為預訓練的擴散模型提供了一種確定性加速采樣方法,該方法以犧牲多樣性來生成質量更好的圖像樣本。隱式采樣利用了廣義的非馬爾可夫正向過程公式:

(17)

將數據分布按照特殊選擇的標準差λ進行重寫:

(18)

用方差表示其均值:

(19)

(20)

通過加速采樣的方式,從完整的{1,2,…,T}時間步長中抽取出子序列{τ1,τ2,…,τS},τ1處于逆擴散的最后一步,從而大幅度減少采樣步驟。

τi=(i-1)·T/S+1

(21)

在與隱式采樣保持兼容的同時,將帶有條件y的逆擴散過程納入公式:

(22)

從xT～N(0,I)出發,加入條件y,沿著一條確定性的反向路徑,最終抵達x0。擴散模型過程如圖1所示。

圖1 擴散模型過程

2 基于分塊的條件擴散圖像恢復

現有的部分生成式網絡模型在進行圖像去雨時要求輸入指定大小的圖片,但是由于數據集的差異,會出現圖片大小不統一的問題。雖然可以通過對圖片進行resize來解決,但這種方法會導致原始信息的丟失和圖片精度的損失。相比之下,該文采用基于SR3[15]的全卷積網絡架構,它使用了來自文獻[16]的U-Net結構的變體,并用BigGAN[18]的殘差塊替換了傳統的殘差塊。為保證模型不受圖像分辨率的影響[19],去掉了自注意力機制、位置編碼和群組歸一化,使得條件擴散模型能夠支持任意圖像大小的輸入。

該文采用圖像重疊分塊處理與重構的方法[20],該方法將圖像分成小塊進行局部操作,并在優化后將結果合并,如圖2所示。在圖像分塊時,保證相鄰圖像塊有一定的重疊,以消除邊縫,并在采樣過程中對重疊小塊之間進行平滑處理。使用圖像分塊處理不僅可以實現并行處理,而且還可以將整張圖像分塊處理并調入內存,提高算法的適用性。這種方法可以避免處理過大圖像時出現的問題,減少每次運算的數據量和計算壓力。

圖2 重疊分塊

為了訓練模型,該文使用帶參數θ的神經網絡進行如下近似:

(23)

(24)

(25)

圖3 基于條件擴散隱式模型圖像去雨

隨后將重疊小塊的噪聲估計累積到矩陣At中,并使用平滑的全圖像噪聲估計執行隱式采樣更新:

(26)

3 實驗結果與分析

3.1 實驗設置

實驗環境為ubuntu18.04,32 GB內存,5 vCPU Intel(R) Xeon(R) Silver 4210 CPU @ 2.20 GHz,RTX 3090(24 GB)GPU,Cuda為11.1,使用Python 3.8編程語言和Pytorch 1.8.1框架進行訓練。在訓練過程中,模型將圖片隨機裁剪為64×64的小塊,并使用Adam優化器,學習率固定為0.000 02,為促進更穩定的學習[17],參數更新時滑動平均值設置為0.999。

3.2 數據集與評價指標

在訓練和測試的過程中,通常需要大量的數據集,尤其是成對的有雨和無雨圖像。為解決這些問題,該文選用公開的訓練集進行訓練和測試,包括Rain14000[21]和SPA-Data[22]。為了評估算法的去雨效果,采用峰值信噪比(PSNR)和結構相似性(SSIM)兩個指標。實驗數據集的劃分如表1所示。

表1 實驗數據集的劃分

3.3 合成雨圖實驗對比

將文中算法與DID-MDN[8],PReNet[9],MPRNet[10],MOSS[11]去雨算法在合成數據集Rain14000上進行定量分析,對比結果如表2所示。與其他算法相比,文中算法在峰值信噪比和結構相似性方面表現更好。同時,圖4展示了文中算法在合成數據集上去雨的視覺效果。從圖中可以發現,PReNet在密集雨線下會有明顯的雨痕殘留,而DID-MDN產生了較大的偽影,MPRNet和MOSS丟失了邊緣紋理細節較差,與之相比,文中算法可以有效地去除密度較大、方向不同的雨紋,并保留更多的圖像細節和邊緣信息,從視覺上得到更令人滿意的效果。

表2 不同算法在合成數據集上的定量對比結果

圖4 不同算法在合成數據集上去雨結果的視覺對比

3.4 真實雨圖實驗對比

為了進一步驗證算法在真實環境下去雨的泛化性能,測試采用了SPA-Data數據集,測試結果如圖5所示。從圖中可以發現,對比算法去除雨滴和較粗的雨紋方面效果并不理想,相比之下,文中算法在雨紋去除和保持圖像細節兩方面都表現最佳,因此在自然雨圖上具有更好的泛化能力。

圖5 不同算法在真實數據集上去雨結果的視覺對比

4 結束語

該文提出一種基于條件擴散隱式模型的圖像去雨方法,采用基于SR3[5]的U-Net全卷積網絡架構,支持任意大小的圖像輸入。引入確定性加速采樣來提高圖像恢復效率。此外,采用圖像重疊分塊處理以減少內存壓力,提高算法的適用性。由實驗分析結果所示,該方法在單圖像去雨效果上有更好的性能,但依然存在不足之處。盡管使用了加速確定性采樣技術,但相較于其他生成模型,圖像恢復時間仍然有待提高,這也是未來需要重點研究的方向。