AL-GAN:一種融合注意力機制的輕量級GAN水下圖像增強模型

馮建新,韓亞軍,潘成勝,孫傳林,蔡遠(yuǎn)航

(1.遼寧省通信網(wǎng)絡(luò)與信息處理重點實驗室, 遼寧 大連 116622; 2.大連大學(xué) 信息工程學(xué)院, 遼寧 大連 116622)

0 引言

海洋有著很多重要的信息資源,這些信息資源對研究海洋生物以及勘探資源有著重要的意義,因此安全可靠地獲取海洋信息資源非常重要。然而海洋環(huán)境不同于陸地,水下環(huán)境情況復(fù)雜,且可見光在水下會遇到吸收和散射等問題,這些會造成采集到的水下信息(如視頻或圖像)會出現(xiàn)偏色、細(xì)節(jié)模糊等問題,這類問題會對后續(xù)水下目標(biāo)跟蹤和檢測造成嚴(yán)重的影響。因此需要提高水下圖像質(zhì)量,消除偏色、細(xì)節(jié)模糊等問題。

用于圖像增強的方法根據(jù)成像模型不同可以分為物理模型和非物理模型2種。水下圖像增強的物理模型是根據(jù)光學(xué)原理建模,主要進(jìn)行去噪和顏色校正,如暗通道先驗(DCP)[1],水下DCP等。水下圖像的精確恢復(fù)取決于模型的先驗知識,然而先驗知識在水下場景不同時往往魯棒性較弱,從而會導(dǎo)致嚴(yán)重的估計偏差。同時,在實際應(yīng)用中水下深度和光傳播系數(shù)基本參數(shù)也很難獲得。水下圖像增強的非物理模型不考慮水下圖像的物理退化機制,它是通過修改退化圖像的具體像素值增強圖像效果,主要增強方法如直方圖滑動拉伸[2]、多尺度融合[3]等。此類方法不依賴于物理成像模型,因此這些方法通常不足以恢復(fù)原始場景特征,尤其是顏色特征。

由于深度學(xué)習(xí)可以實現(xiàn)復(fù)雜非線性系統(tǒng)的端到端建模,能夠提高感知圖像的質(zhì)量,目前已經(jīng)在去噪、去模糊、感知增強和對比度調(diào)整等低級視覺任務(wù)中取得了令人信服的成功[4-5]。尤其是近些年生成對抗網(wǎng)絡(luò)(generative adversarial networks,GAN)[6]發(fā)展迅速,因此越來越多的研究人員將目光轉(zhuǎn)移到了應(yīng)用深度學(xué)習(xí)進(jìn)行水下圖像增強的研究中。JieLi等[7]提出了一種無監(jiān)督生成對抗網(wǎng)絡(luò)Water-GAN實現(xiàn)對單目圖像水下偏色問題的矯正,它具有很好的實時性。Fabbri等[8]用擴(kuò)充數(shù)據(jù)集方法提出適用于水下場景的UGAN(underwater GAN),解決了水下圖像數(shù)據(jù)集缺乏問題,并且能夠有效地提升水下圖像的對比度和清晰度。Islam等[9]提出了一種基于GAN實時性很高的水下圖像增強模型FUnIE-GAN,該模型根據(jù)圖像的整體內(nèi)容、細(xì)節(jié)樣式、風(fēng)格和局部紋理信息建立損失函數(shù),可以實時增強水下圖像,網(wǎng)絡(luò)模型推理速度較快,但是針對不同的水體情況處理效果不好,例如過度增強或增強不足,同時模型魯棒性會降低。Zhang等[10]改進(jìn)原有生成對抗網(wǎng)絡(luò),為了避免圖像細(xì)節(jié)模糊,通過引入梯度損失銳化圖像邊緣提高圖像清晰度,但是在處理偏色問題上會有顏色失真問題。Hambarde等[11]提出一種端到端的水下生成對抗網(wǎng)絡(luò)UW-GAN,能夠完成單個水下圖像深度預(yù)測,進(jìn)一步提出了一種用于大規(guī)模數(shù)據(jù)庫合成水下圖像生成方法。UW-GAN網(wǎng)絡(luò)模型在U-Net框架生成網(wǎng)絡(luò)中使用VGG16模型作為編碼器,雖然圖像增強效果較好,但模型參數(shù)量較大,網(wǎng)絡(luò)模型的推理速度很慢。

上述基于GAN的模型直接學(xué)習(xí)退化圖像到清晰圖像的映射,針對不同水體情況無法做到增強水下圖像質(zhì)量的同時提高網(wǎng)絡(luò)模型的推理速度。為了做到增強不同水體的水下圖像質(zhì)量,兼顧網(wǎng)絡(luò)模型推理速度的提升,本文中提出了一種融合注意力機制的輕量級GAN模型,以下是本文中的主要研究內(nèi)容:

1) 提出了一種輕量級GAN模型,生成網(wǎng)絡(luò)在FUnIE-GAN模型基礎(chǔ)上使用MobileNet替換原U-Net生成網(wǎng)絡(luò)參數(shù)量極大的VGG16模型,減少了模型參數(shù)量,從而解決生成網(wǎng)絡(luò)推理速率慢的問題。與其他模型比較推理速度提高了1.2～3倍。

2) 在特征提取網(wǎng)絡(luò)中引入了通道和空間雙重注意力機制,實現(xiàn)了去除水下圖像顏色干擾和增強圖像細(xì)節(jié)的目的。

3) 所提出的模型在數(shù)據(jù)集上表現(xiàn)出更好、更穩(wěn)定的性能。在EUVP數(shù)據(jù)集上,與其他算法相比,水下圖像評價指標(biāo)(underwater image quality measure,UIQM)[12]提高約0.21,自然圖像評價指標(biāo)(natural image quality evaluator,NIQE)[13]降低約0.65。

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

網(wǎng)絡(luò)模型包括生成網(wǎng)絡(luò)和判別網(wǎng)絡(luò)2部分,如圖1所示。本文中生成網(wǎng)絡(luò)的輸入為水下退化圖像,判別網(wǎng)絡(luò)的輸入為生成圖像和地面真實圖像。在模型訓(xùn)練階段,生成網(wǎng)絡(luò)和判別網(wǎng)絡(luò)互相監(jiān)督,循環(huán)交替訓(xùn)練更新網(wǎng)絡(luò)模型的參數(shù)。最終訓(xùn)練出高性能并且穩(wěn)定的生成網(wǎng)絡(luò),用生成網(wǎng)絡(luò)來增強水下圖像。

U-Net[14]是一個包含上采樣和下采樣的全卷積網(wǎng)絡(luò),它通過在上采樣過程中組合不同下采樣的特征重新構(gòu)建清晰圖像。我們的生成網(wǎng)絡(luò)受U-Net結(jié)構(gòu)啟發(fā),設(shè)計為5層編解碼結(jié)構(gòu),編碼器模塊由深度可分離卷積模塊和注意力模塊組成,解碼器模塊由多個上采樣卷積模塊組合而成。為了減少網(wǎng)絡(luò)在下采樣期間造成的信息丟失,U型網(wǎng)絡(luò)連接編碼器輸出作為解碼器輸入。判別網(wǎng)絡(luò)模型通過定義PatchGAN[15]實現(xiàn)。

1.1 生成網(wǎng)絡(luò)模型

1.1.1編碼器模塊

根據(jù)MobileNet[16]模型體積小、精度高特性,深度可分離卷積作為MobileNet的核心組件,由深度卷積和逐點卷積組成。在編碼器模塊中使用深度可分離卷積替換普通卷積,深度可分離卷積用于特征提取的卷積大小為3×3,用來調(diào)整通道數(shù)的卷積塊大小為1×1。

CBAM[17]注意力模塊由通道注意模塊(CAM)和空間注意模塊(SAM)組成。輸入圖像的特征首先通過卷積模塊提取出來,然后再輸入到CAM得到輸入特征的權(quán)重分布圖,權(quán)重分布圖可以顯示輸入圖像信息的重要度。SAM側(cè)重于輸入特征信息在哪里更重要。2種注意機制結(jié)合增強了網(wǎng)絡(luò)的學(xué)習(xí)功能,使網(wǎng)絡(luò)更加注重水下圖像退化嚴(yán)重部分的信息,從而能夠更好地增強圖像,本文中CBAM結(jié)構(gòu)如圖2所示。

編碼器模塊經(jīng)過前3層卷積模塊后接入CBAM注意力模塊,卷積模塊的每一層先是使用大小3×3,步長為2的深度卷積提取特征,然后用1×1卷積調(diào)整通道數(shù),每個卷積層后面連接的是歸一化BN及ReLU激活函數(shù)。為探尋注意力機制和深度可分離卷積模塊的最優(yōu)組合,在3.2節(jié)進(jìn)行了生成網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計實驗,實驗結(jié)果表明,前3層卷積模塊和CBAM組合結(jié)構(gòu)是最優(yōu)組合。

1.1.2解碼器模塊

解碼器由5個上采樣卷積模塊組成,解碼器中前4個卷積模塊的每一層采用大小為3×3的濾波器,卷積步長為2。在每個卷積層之后添加BN和激活函數(shù)ReLU,最后一個卷積模塊將特征圖轉(zhuǎn)換為256×256×3輸出圖像。

圖1 網(wǎng)絡(luò)模型結(jié)構(gòu)

圖2 通道空間注意力模塊

1.2 判別網(wǎng)絡(luò)模型

本文中判別網(wǎng)絡(luò)使用馬爾可夫鑒別器(PatchGAN),與常規(guī)鑒別器不同,PatchGAN輸入是當(dāng)前輸入圖像Ii的地面真值圖像IGT和其經(jīng)過生成網(wǎng)絡(luò)處理得到的輸出圖像IG,輸出大小為16×16特征矩陣,這在判斷之前就將圖像進(jìn)行了小塊分割,從而提高了鑒別器精度,同時也促使網(wǎng)絡(luò)在訓(xùn)練過程中可以獲得更多的細(xì)節(jié)信息。本文中判別網(wǎng)絡(luò)前4層使用3×3的卷積層并進(jìn)行2倍下采樣,在每一卷積層后添加BN層和Leaky ReLU激活層,最后一層在卷積層后添加Tanh激活層,5個卷積層將256×256×6的輸入圖像Ii(生成圖像和真實圖像)轉(zhuǎn)換為16×16×1的輸出,輸出矩陣的每一個元素都表示輸入圖像的一個感受野,從而可以更好地捕獲圖像局部特征,如局部紋理和細(xì)節(jié)形式[18],判別網(wǎng)絡(luò)結(jié)構(gòu)如圖3所示。

圖3 判別網(wǎng)絡(luò)結(jié)構(gòu)

1.3 損失函數(shù)

為了更好恢復(fù)圖像視覺效果,同時盡可能保留細(xì)節(jié)特征,損失函數(shù)包括對抗損失、全局相似性損失、內(nèi)容損失3個方面,損失函數(shù)定義如下:

(1)

式(1)中:LWGAN[19]為對抗損失函數(shù);LL1[20]為全局相似性損失函數(shù);Lcon[21]為內(nèi)容損失函數(shù);λ1、λ2為平衡各項損失函數(shù)的權(quán)重因子,其中λ1=0.7,λ2=0.3。G和D分別表示生成網(wǎng)絡(luò)和判決網(wǎng)絡(luò)。

1.3.1對抗損失LWGAN

傳統(tǒng)GAN損失函數(shù)是基于散度做優(yōu)化,但是在這種優(yōu)化下網(wǎng)絡(luò)損失訓(xùn)練會很不穩(wěn)定,會帶來梯度消失等問題,這種問題會導(dǎo)致模型訓(xùn)練崩潰[22]。因此,使用帶有梯度懲罰的WGAN(WGAN-GP)[23],對抗損失函數(shù)表示為

LWGAN(G,D)=E[D(y)]-E[D(G(x))]+

(2)

1.3.2全局相似性損失LL1

由于L1損失不容易引入模糊,所以使用輸出圖像Io與生成網(wǎng)絡(luò)輸出圖像IG之間L1距離優(yōu)化鑒別器,從而保證圖像信息的一致性,其定義如下:

(3)

1.3.3內(nèi)容損失Lcon

為使生成圖像更加真實,提升增強圖像的視覺效果,在目標(biāo)函數(shù)中添加內(nèi)容損失,本文中使用感知距離作為內(nèi)容損失形式[24]。引入VGG19預(yù)訓(xùn)練網(wǎng)絡(luò),通過提取第4次最大池化前第3個卷積層輸出的高級特征構(gòu)造內(nèi)容損失,定義如下,其中φ代表網(wǎng)絡(luò)提取到的高級特征。

(4)

2 實驗設(shè)計

2.1 實驗參數(shù)

本文中實驗是基于Pytorch深度學(xué)習(xí)框架實現(xiàn)的,實驗CPU為12th Gen Intel(R) Core(TM) i5-12500H 3.10 GHz,GPU為NVIDIA GeForce RTX3050Ti。訓(xùn)練模型的優(yōu)化器采用Adam[25],批處理大小Batch-size為8,初始學(xué)習(xí)率設(shè)為0.01,一共150個epoch,每30個epoch學(xué)習(xí)率乘以0.5衰減1次。全局相似性損失權(quán)重系數(shù)λ1為0.7,內(nèi)容損失權(quán)重系數(shù)λ2為0.3。

2.2 生成網(wǎng)絡(luò)結(jié)構(gòu)

從訓(xùn)練過程和訓(xùn)練結(jié)果出發(fā),對生成網(wǎng)絡(luò)U-Net框架編碼器部分CBAM模塊和卷積模塊進(jìn)行組合訓(xùn)練測試。選擇UIQM和網(wǎng)絡(luò)模型推理速度作為衡量網(wǎng)絡(luò)性能指標(biāo)。針對不同的CBAM模塊與下采樣卷積模塊組合,表1是采用EUVP數(shù)據(jù)集中相同測試集進(jìn)行測試得到的性能,其中F1表示第1層卷積模塊與CBAM模塊組合,F2表示第2層卷積模塊與CBAM模塊組合,F3和F4同理。

從表1可以看出,第2層卷積模塊與CBAM模塊組合方案F4在模型推理速度上與F1相當(dāng),略高于F2,遠(yuǎn)低于F3,但其UIQM為2.327,說明F4在處理圖像效果上明顯優(yōu)于其他組合。

表1 單層卷積模塊和CBAM組合結(jié)果對比

表2為2層上采樣卷積模塊和CBAM組合的測試結(jié)果。從表2可以看出,2層卷積模塊與CBAM組合方案明顯比單層卷積模塊與CBAM結(jié)合方案效果好,其中F2+F3的UIQM可達(dá)2.465,但是模型推理速度相對變慢,F2+F4為1.73 s。

表2 兩層卷積模塊和CBAM組合結(jié)果對比

表3為3/4層上采樣卷積模塊和CBAM組合方案的測試結(jié)果。從表3可以看出,3/4層卷積模塊和CBAM組合方案的UIQM比2層組合更優(yōu),但同時模型推理速度也變慢。3層卷積模塊與CBAM組合方案和每一層卷積模塊都與CBAM組合方案相比,UIQM值幾乎一致,其中F1+F2+F3方案推理速度為1.75 s最快,該值與2層采樣卷積模塊和CBAM組合方案相近。

綜上所述,從推理速度和圖像效果出發(fā),本文中生成網(wǎng)絡(luò)確定采用F1+F2+F3方案,即前3層卷積模塊與CBAM組合方案構(gòu)造生成網(wǎng)絡(luò)的編碼器。

2.3 損失函數(shù)性能

圖4和圖5是在對抗損失、全局相似性損失、內(nèi)容損失3個方面構(gòu)造損失函數(shù)的性能分析。其中,Lall為3種損失函數(shù)組合,LWGAN+L1為對抗損失函數(shù)與全局相似性損失函數(shù)組合,LWGAN+C為對抗損失函數(shù)與內(nèi)容損失函數(shù)的組合;n為迭代次數(shù),loss為損失函數(shù)值。

從圖4可以看出,只使用對抗性損失函數(shù)訓(xùn)練時震動幅度大,收斂過程不穩(wěn)定,所需要循環(huán)次數(shù)多。采用3種損失函數(shù)組合形式的訓(xùn)練過程更加平穩(wěn),離散程度更小。

從圖5可以看出,Lall曲線、LWGAN+L1曲線表現(xiàn)相近且損失函數(shù)值明顯低于LWGAN+C曲線。但隨著迭代次數(shù)增加,LWGAN+L1曲線收斂后又出現(xiàn)了小幅度的波動,Lall曲線相比于LWGAN+L1曲線更加穩(wěn)定。綜上所述,使用3種損失函數(shù)組合Lall具有更好、更穩(wěn)定的性能。此外,從峰值信噪比(PSNR)和結(jié)構(gòu)相似性(SSIM)[26]可知,3種損失函數(shù)組合Lall的PSNR是23.749,SSIM是0.765,明顯高于其他損失函數(shù)組合,因此,本文中損失函數(shù)考慮對抗損失、全局相似性損失、內(nèi)容損失3個方面,損失函數(shù)為3種損失函數(shù)組合Lall,如式(1)所示。

圖5 3種損失函數(shù)訓(xùn)練過程對比

表4為不同損失函數(shù)組合訓(xùn)練出的模型在評價指標(biāo)上的實驗結(jié)果對比。

表4 不同損失函數(shù)組合評價指標(biāo)實驗結(jié)果對比

3 結(jié)果與討論

選用水下公開數(shù)據(jù)集EUVP驗證本文中算法的有效性,EUVP數(shù)據(jù)集包含12 000多對水下圖像和其對應(yīng)的地面清晰圖像,選取其中7 000幅水下圖像和其對應(yīng)的清晰圖像作為訓(xùn)練集。將本文中算法和傳統(tǒng)算法、深度學(xué)習(xí)算法從主客觀兩方面做對比,由實驗結(jié)果可知本文中算法在復(fù)雜情況下的水下圖像增強效果更好。

3.1 主觀評價

為了驗證本文中算法對水下真實圖像增強效果,測試數(shù)據(jù)隨機選取150張真實水下圖像。圖6為測試數(shù)據(jù)部分圖像在不同算法上的效果圖,對比的傳統(tǒng)算法包括水下暗通道先驗(UDCP),基于圖像模糊的水下圖像恢復(fù)(IBLA)[27],基于快速場景深度估計模型的水下圖像恢復(fù)(ULAP)[28];對比的深度學(xué)習(xí)算法包括基于對抗的水下圖像增強網(wǎng)絡(luò)(UGAN),基于殘差網(wǎng)絡(luò)的生成模型(DeepSESR)[29],基于GAN網(wǎng)絡(luò)的快速水下圖像增強(FUnIE-GAN)。圖6(a)為輸入圖像,圖6(I)為真實水下圖像,圖6(b)—圖6(h)分別是UDCP、IBLA、ULAP、UGAN、Deep-SESR、FUnIE-GAN以及本文中方法的處理結(jié)果圖。可見本文中方法相比于其他方法結(jié)果更加貼近真實圖像,解決了水下圖像由于水體原因造成的偏色問題,處理得到的圖像有更高的對比度和清晰度,色彩更加鮮艷。

圖6 不同方法處理結(jié)果對比

3.2 客觀評價

客觀地分析和評價算法性能,使用水下圖像質(zhì)量指標(biāo)UIQM和自然圖像評價指標(biāo)NIQE作為評價指標(biāo)。UIQM是由水下圖像色彩測量UICM,水下圖像清晰度測量UISM,水下圖像對比度測量UIconM,3個評價指標(biāo)組成的。UIQM的值越大,表明圖像的質(zhì)量越高,其計算公式為

UIQM=c1×UICM+c2×UISM+c3×UIconM

(5)

根據(jù)文獻(xiàn)[30]設(shè)置c1=0.028 2,c2=0.295,c3=3.575。

NIQE無須利用人眼評分的失真圖像進(jìn)行訓(xùn)練,NIQE相比于傳統(tǒng)評價指標(biāo)PNSR和SSIM能更有效地反映圖像重建的質(zhì)量。NIQE指標(biāo)的值越小,表明圖像的質(zhì)量越高,其計算公式為

D(v1,v2,∑1,∑2)=

(6)

式(6)中:v1和v2分別表示自然多元高斯模型和失真圖像多元高斯模型的均值向量;∑1和∑2分別表示自然MVG模型和失真圖像MVG模型的協(xié)方差矩陣。

對150張水下圖像進(jìn)行評估,表5為各種方法處理圖像的評價指標(biāo)UIQM和NIQE。可以看出,經(jīng)本文中算法處理的圖像評價指標(biāo)最優(yōu)。與6種對比算法相比,本文中算法處理的圖像UIQM明顯得到提升,同時NIQE明顯降低。與FUnIE-GAN算法相比,UIQM提升了約0.21 ,NIQE降低了約0.65。

表5 UIQM和NIQE的定量比較

綜上所述,根據(jù)各種算法在無參考圖像質(zhì)量評價指標(biāo)UIQM和NIQE上的對比,證明本文中算法能夠更好的恢復(fù)水下圖像,提高了圖像的清晰度和對比度。

3.3 消融實驗

表6為采用VGG16模型做編碼器和本文中編碼器在處理單幅圖像所需時間。隨機選取5張圖片,分別在VGG16和MobileNet為編碼器的網(wǎng)絡(luò)中測試,結(jié)果表明,本文中編碼器比VGG16編碼器運算速度提高很大,表明MobileNet替換VGG16的有效性。

表6 不同編碼器運行時間對比(s)

圖7為有無CBAM模塊處的圖像對比,從圖7前2列對比圖可以看出,無CBAM模塊處理的圖像較有CBAM模塊處理的圖像雖然解決了水下圖像偏綠色問題,但是增強后的圖像明顯出現(xiàn)偏黃。從后面3列對比圖可以看出,無CBAM模塊處理的圖像在細(xì)節(jié)上雖然得到增強,但是相比于有CBAM模塊處理仍然有模糊問題。有CBAM模塊處理的圖像,不僅解決了色偏問題,而且清晰度也得到提升,因此本文中模型引入了CBAM模塊。

圖7 有無通道空間注意力模塊結(jié)果對比

表7用客觀指標(biāo)UIQM和NIQE進(jìn)行了分析,實驗結(jié)果如表7所示。從表7可見,有CBAM模塊處理的圖像UIQM明顯比無CBAM模塊處理的高,同時NIQE明顯更低,定量說明有CBAM模塊處理的圖像質(zhì)量更高。

表7 有無CBAM模塊UIQM/NIQE評價指標(biāo)結(jié)果對比 Table 7 Comparison of UIQM/NIQE evaluation index results with or without CBAM module

3.4 運行時間

進(jìn)一步,為驗證本文中提出的GAN水下圖像增強模型輕量化特征,對方法運行時間進(jìn)行了估計處理,其中Deep-SESR、UGAN和FUnIE-GAN是基于深度學(xué)習(xí)的方法,結(jié)果如表8所示。該結(jié)果為CPU處理1張大小256×256圖像所需要的時間,可以看出本文中方法雖然不是最快的,但是處理時間與FUnIE-GAN相當(dāng)。

4 結(jié)論

本文中提出一種AL-GAN水下圖像增強模型。實驗結(jié)果表明,相比與對比算法,本文中提出的水下圖像增強模型UIQM提高了約0.21,NIQE降低了約0.65,說明本文中模型能夠更好地增強圖像細(xì)節(jié),提升圖像的對比度和清晰度,對于水下圖像增強效果表現(xiàn)最優(yōu),同時網(wǎng)絡(luò)模型的推理速度同比于其他模型提高了1.2～3倍。

但在不同的水下環(huán)境中,水下圖像退化的程度也會有所不同。嚴(yán)重退化的圖像單獨使用深度學(xué)習(xí)網(wǎng)絡(luò)處理后,仍然會有嚴(yán)重的顏色干擾,并可能產(chǎn)生冗余噪聲點。對深度學(xué)習(xí)網(wǎng)絡(luò)添加物理模型可以更好地估計水下圖像成像環(huán)境,進(jìn)一步,通過環(huán)境參數(shù)約束網(wǎng)絡(luò)可以更好地恢復(fù)出清晰的水下場景圖像。在未來的工作中,我們將嘗試結(jié)合物理模型,進(jìn)一步關(guān)注退化嚴(yán)重的水下圖像處理。