基于擴(kuò)散模型的圖像去噪方法研究

薛永航，白 帆，李 娜

（沈陽理工大學(xué)，遼寧 沈陽 110168）

0 引言

在實(shí)際工程中，圖像采集和傳輸過程中會(huì)受到外界因素干擾，尤其噪聲的影響，從而導(dǎo)致圖像特征提取不準(zhǔn)確。因此，如何能夠有效去除圖像中的噪聲，在圖像處理領(lǐng)域具有重要意義。傳統(tǒng)去噪方法容易造成信息丟失、計(jì)算復(fù)雜度高和參數(shù)選擇困難等問題。隨著深度學(xué)習(xí)的發(fā)展[1]，從卷積神經(jīng)網(wǎng)絡(luò)、循環(huán)神經(jīng)網(wǎng)絡(luò)、到變分神經(jīng)網(wǎng)絡(luò)。這些方法從圖像特征角度處理噪聲，難以學(xué)習(xí)時(shí)頻圖像噪聲分布，由此導(dǎo)致時(shí)頻圖噪聲去除時(shí)邊緣信息丟失多，效果較差[2]。為了提高神經(jīng)網(wǎng)絡(luò)在時(shí)頻圖像的噪聲去除效果，通過Wigner-Ville（Wigner-Ville Distibution，WVD）分布算法[3]，提供了時(shí)頻圖像在時(shí)域和頻域上的局部信息，使得在時(shí)間和頻率上同時(shí)具有較高的分辨率[4]。讓網(wǎng)絡(luò)學(xué)習(xí)到噪聲圖像和無噪聲圖像的特征，極大降低了該算法在通用數(shù)據(jù)集上的敏感程度，有效提升了模型泛化能力。

本文通過基于LoRA 的擴(kuò)散模型神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)[5]，將其應(yīng)用于時(shí)頻圖像去噪，解決了在復(fù)雜工況下時(shí)頻圖像噪聲難以去除的問題，通過學(xué)習(xí)含噪信號(hào)及無噪信號(hào)圖像概率分布，達(dá)到去除噪聲的目的。實(shí)驗(yàn)結(jié)果表明，相對(duì)于傳統(tǒng)深度學(xué)習(xí)去噪模型，該模型能夠有效去除時(shí)頻圖像的噪聲。

1 算法原理

1.1 基于LoRA 的去噪擴(kuò)散模型

該算法將擴(kuò)散模型以及LoRA 網(wǎng)絡(luò)結(jié)構(gòu)相結(jié)合，用于時(shí)頻信號(hào)圖像的去噪，整個(gè)網(wǎng)絡(luò)分為兩大部分，訓(xùn)練階段和采樣階段。

（1）訓(xùn)練階段

首先，使用已預(yù)訓(xùn)練的自編碼器將時(shí)頻圖像從像素空間映射到潛在空間，以學(xué)習(xí)時(shí)頻圖像的隱式表達(dá)，同時(shí)壓縮圖像尺寸，減少計(jì)算該網(wǎng)絡(luò)計(jì)算復(fù)雜度和模型的參數(shù)量。然后，由文本編碼器對(duì)含有噪聲的時(shí)頻圖像的提示文本進(jìn)行編碼，生成大小為[B，K，E] 的嵌入文本詞。其中，B 表示批次大小，K 表示文本的最大編碼長度，E 表示嵌入的維度，嵌入文本能夠有效捕捉時(shí)頻圖像語義特征?；贚oRA 的擴(kuò)散模型輸出噪聲?θ，通過計(jì)算時(shí)頻圖像與真實(shí)噪聲之間的誤差作為損失函數(shù)。通過反向傳播算法，更新模型中的參數(shù)。

（2）采樣階段

采樣階段包括對(duì)時(shí)頻圖像文本編碼和時(shí)頻圖像解碼過程。在文本編碼過程中，文本編碼器由輸入的描述文本標(biāo)簽，對(duì)時(shí)頻圖像的描述文本進(jìn)行編碼，包括時(shí)頻噪聲生成與恢復(fù)潛在表示，隨機(jī)產(chǎn)生大小為[B，Z，H/8，W/8] 的噪聲，利用訓(xùn)練好的模型，對(duì)網(wǎng)絡(luò)中參數(shù)進(jìn)行迭代，逐步去除時(shí)頻圖像中的噪聲，恢復(fù)時(shí)頻圖像的潛在表示。

1.2 LoRA 結(jié)構(gòu)應(yīng)用

LoRA 預(yù)訓(xùn)練模型（Pretrianed Weight）可以用于時(shí)頻圖像去噪任務(wù)的模塊中，通過調(diào)整替換矩陣A和B，可以達(dá)到凍結(jié)共享模型的作用，能夠有效的切換不同訓(xùn)練任務(wù)，即訓(xùn)練時(shí)頻圖像和圖像描述文本，從而顯著的降低數(shù)據(jù)參數(shù)存儲(chǔ)需求，使得LoRA 的訓(xùn)練更加有效。

通過訓(xùn)練時(shí)頻圖像的提示文本，可以替代整個(gè)擴(kuò)散模型的參數(shù)訓(xùn)練，在原有對(duì)圖像特征提取的基礎(chǔ)上，可以使該算法識(shí)別到噪聲圖像的準(zhǔn)確率有效提升。

圖1 LoRA 結(jié)構(gòu)

1.3 擴(kuò)散過程

擴(kuò)散過程中將隨機(jī)噪聲添加到數(shù)據(jù)中，從噪聲中構(gòu)造所需的數(shù)據(jù)樣本。與VAE 模型不同，擴(kuò)散模型是通過固定過程學(xué)習(xí)，并且潛在變量具有與原始數(shù)據(jù)相同方法。

（1）前向過程

逐步添加高斯噪聲到圖像中，得到一個(gè)含有噪聲的時(shí)頻圖像。其中，Xt時(shí)刻的分布等于Xt-1時(shí)刻的分布加上高斯分布的噪聲。

式中：Xt表示t時(shí)刻圖像，其中t?[1，T]，αt是噪聲的衰減值，Z表示高斯噪聲。

重復(fù)迭代后，可由初始狀態(tài)X0得任意t時(shí)刻分布：

（2）逆向過程

在此過程中，逐步從噪聲中復(fù)原出原始時(shí)頻圖。由貝葉斯公式可知：

q（Xt-1|Xt）為給定Xt時(shí)Xt-1的概率分布，由于q（Xt-1）分布無法直接求取，因此用條件概率分布q（Xt-1|Xt，X0）近似求解，如（4）式所示。

在已知X0的情況下，可以求出Xt-1時(shí)刻的分布。

重復(fù)迭代后，最終可預(yù)測出X0時(shí)刻的分布。

式中：βt= 1 -αt，?θ（Xt，t）表示t時(shí)刻模型預(yù)測的噪聲，σt表示預(yù)測的噪聲方差。

2 實(shí)驗(yàn)結(jié)果與分析

2.1 實(shí)驗(yàn)流程

為了驗(yàn)證該算法去噪性能，本次研究采用一組工況數(shù)據(jù)集，該數(shù)據(jù)是兩相電流信號(hào)，維度是2×10000，首先，對(duì)該數(shù)據(jù)維度進(jìn)行處理，然后，將二維信號(hào)處理成一維信號(hào)，截取部分信號(hào)數(shù)據(jù)，通過WVD 算法，將該數(shù)據(jù)轉(zhuǎn)成時(shí)頻圖，以便于更好捕捉信號(hào)在不同頻率上的瞬時(shí)特性。并且，對(duì)原始數(shù)據(jù)添加均值為0.2 和標(biāo)準(zhǔn)差為0.35 的高斯噪聲，以驗(yàn)證算法的可靠性。整個(gè)訓(xùn)練過程，選擇Adam（Adaptive Moment Estimation）參數(shù)優(yōu)化器，學(xué)習(xí)率設(shè)置為0.001，訓(xùn)練150 個(gè)周期。

2.2 評(píng)價(jià)指標(biāo)

（1）峰值信噪比計(jì)算（Peak signal-to-noise ratio，PSNR）：

峰值信噪比通過均方差（MSE）進(jìn)行定義，方差定義為：

由方差可知，峰值信噪計(jì)算公式如下：

（2）結(jié)構(gòu)相似度計(jì)算（Structural Similarity，SSIM）：

通過從圖像的亮度、對(duì)比度、結(jié)構(gòu)三個(gè)方面度量圖像相似性，計(jì)算如下：

其中，μX，μY分別表示圖像X和Y的均值，σX，σY分別表示圖像X和Y的方差，σX σY表示圖像X和Y的協(xié)方差，C1、C2、C3表示一個(gè)常數(shù)。SSIM 的取值范圍在[0，1]之間，它的值越大，表示圖像的失真越小。

2.3 實(shí)驗(yàn)結(jié)果

通過WVD 算法處理后的時(shí)頻圖像，數(shù)據(jù)信息被壓縮，時(shí)頻圖像特征信息區(qū)分不明顯，在VAE 算法中，首先對(duì)原始時(shí)頻圖圖像進(jìn)行編碼，輸入的時(shí)頻圖像X 通過編碼器輸出兩個(gè)M 維向量，這兩個(gè)向量是潛在空間Z 的兩個(gè)參數(shù)。其次，在潛在空間Z 中增加約束條件，這個(gè)約束條件迫使?jié)撛诳臻gZ 產(chǎn)生服從單位正態(tài)分布的潛在變量。最后，Z 通過解碼器生成一個(gè)樣本，由于是隨機(jī)采樣，從而導(dǎo)致潛在空間Z 的不確定性，變分后驗(yàn)難以選擇。因此，當(dāng)時(shí)頻圖像中噪聲分布不均勻時(shí)，實(shí)際輸出時(shí)頻圖像數(shù)據(jù)和輸入圖像數(shù)據(jù)出現(xiàn)較大偏差，導(dǎo)致VAE 難以學(xué)習(xí)原時(shí)頻圖像分布，無法近似真實(shí)后驗(yàn)，網(wǎng)絡(luò)模型去噪效果差，且容易引入新的噪聲。由圖2 可知，a 圖噪聲去除但時(shí)頻信息損失較多，b 圖噪聲仍分布于中間區(qū)域，未去除完全。圖像的峰值信噪難以達(dá)到20，其結(jié)構(gòu)相似度低于0.9。

相較于VAE 算法，采用基于LoRA 的擴(kuò)散模型去噪方法，先對(duì)隨機(jī)噪聲樣本進(jìn)行采樣，通過逐步采樣噪聲樣本，將采樣后的噪聲樣本進(jìn)行嵌入，包括時(shí)間嵌入以及文本嵌入，時(shí)間嵌入即上文提到的擴(kuò)散過程，通過前向過程和逆向過程，對(duì)時(shí)頻圖像進(jìn)行加噪和去噪，從而能夠有效還原時(shí)頻圖像的峰值信息，其次，文本嵌入引入LoRA 語言模型，極大增強(qiáng)了網(wǎng)絡(luò)對(duì)時(shí)頻圖像噪聲的理解，彌補(bǔ)了網(wǎng)絡(luò)在時(shí)頻圖像的低峰值區(qū)域的不敏感程度，使得該網(wǎng)絡(luò)能夠有效理解時(shí)頻圖像輸入信息以及噪聲水平，該方法在時(shí)頻圖像去噪中效果更好，通過將去噪圖像與原始噪聲圖像對(duì)比分析可知，該時(shí)頻圖像細(xì)節(jié)恢復(fù)較高，噪聲去除效果顯著。不同于VAE 算法中高斯采樣，基于LoRA 的擴(kuò)散模型去噪方法采樣方式更多，有Euler 一階采樣器以及擴(kuò)散概率模型求解器（Dirichlet Process Mixture，DPM），數(shù)據(jù)表明DPM++采樣方式效果最優(yōu)，該算法相較于VAE，峰值信噪比提升20.2%，結(jié)構(gòu)相似度提升11%，訓(xùn)練時(shí)間縮短0.7 h。由圖3 可知，在a 圖及b圖中，該算法針對(duì)時(shí)頻信號(hào)圖像，去除了大部分噪聲并保留了原始圖像的邊緣信息。VAE 去噪算法以及基于LoRA 擴(kuò)散模型去噪算法，計(jì)算峰值信噪比和結(jié)構(gòu)相似度指標(biāo)見表1。該算法去噪效果見圖2 和圖3。

表1 VAE 及擴(kuò)散模型在時(shí)頻信號(hào)圖中PSNR 和SSIM 指標(biāo)

圖3 基于LoRA 的擴(kuò)散模型去噪效果

3 結(jié)語

針對(duì)時(shí)頻圖像的噪聲分布，本研究采用基于LoRA 微調(diào)的擴(kuò)散模型對(duì)時(shí)頻圖像去噪，在時(shí)頻圖像恢復(fù)中取得了顯著的效果。該模型通過深度神經(jīng)網(wǎng)絡(luò)模型，學(xué)習(xí)含有噪聲的時(shí)頻圖像分布，成功地降低了噪聲對(duì)時(shí)頻圖像的影響，并恢復(fù)了原始時(shí)頻圖像的細(xì)節(jié)，提高了時(shí)頻圖像的清晰度。

該算法在峰值信噪比和結(jié)構(gòu)相似度評(píng)價(jià)指標(biāo)上表現(xiàn)出較為理想的結(jié)果，驗(yàn)證了在時(shí)頻圖像去噪任務(wù)上的有效性和可靠性。通過對(duì)比原始噪聲圖像和去噪圖像，可以明顯看出時(shí)頻圖像噪聲顯著減少，去噪質(zhì)量顯著提高。在未來研究中可以考慮優(yōu)化網(wǎng)絡(luò)結(jié)構(gòu)、增加訓(xùn)練數(shù)據(jù)量、引入更多微調(diào)方式進(jìn)行模型訓(xùn)練，以進(jìn)一步提升去噪效果。