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

馮建新,韓亞軍,潘成勝,孫傳林,蔡遠航

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

0 引言

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

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

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

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

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

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

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

1 網絡結構

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

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

1.1 生成網絡模型

1.1.1編碼器模塊

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

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

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

1.1.2解碼器模塊

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

圖1 網絡模型結構

圖2 通道空間注意力模塊

1.2 判別網絡模型

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

圖3 判別網絡結構

1.3 損失函數

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

(1)

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

1.3.1對抗損失LWGAN

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

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

(2)

1.3.2全局相似性損失LL1

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

(3)

1.3.3內容損失Lcon

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

(4)

2 實驗設計

2.1 實驗參數

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

2.2 生成網絡結構

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

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

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

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

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

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

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

2.3 損失函數性能

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

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

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

圖5 3種損失函數訓練過程對比

表4為不同損失函數組合訓練出的模型在評價指標上的實驗結果對比。

表4 不同損失函數組合評價指標實驗結果對比

3 結果與討論

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

3.1 主觀評價

為了驗證本文中算法對水下真實圖像增強效果,測試數據隨機選取150張真實水下圖像。圖6為測試數據部分圖像在不同算法上的效果圖,對比的傳統(tǒng)算法包括水下暗通道先驗(UDCP),基于圖像模糊的水下圖像恢復(IBLA)[27],基于快速場景深度估計模型的水下圖像恢復(ULAP)[28];對比的深度學習算法包括基于對抗的水下圖像增強網絡(UGAN),基于殘差網絡的生成模型(DeepSESR)[29],基于GAN網絡的快速水下圖像增強(FUnIE-GAN)。圖6(a)為輸入圖像,圖6(I)為真實水下圖像,圖6(b)—圖6(h)分別是UDCP、IBLA、ULAP、UGAN、Deep-SESR、FUnIE-GAN以及本文中方法的處理結果圖?？梢姳疚闹蟹椒ㄏ啾扔谄渌椒ńY果更加貼近真實圖像,解決了水下圖像由于水體原因造成的偏色問題,處理得到的圖像有更高的對比度和清晰度,色彩更加鮮艷。

圖6 不同方法處理結果對比

3.2 客觀評價

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

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

(5)

根據文獻[30]設置c1=0.028 2,c2=0.295,c3=3.575。

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

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

(6)

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

對150張水下圖像進行評估,表5為各種方法處理圖像的評價指標UIQM和NIQE?？梢钥闯?經本文中算法處理的圖像評價指標最優(yōu)。與6種對比算法相比,本文中算法處理的圖像UIQM明顯得到提升,同時NIQE明顯降低。與FUnIE-GAN算法相比,UIQM提升了約0.21 ,NIQE降低了約0.65。

表5 UIQM和NIQE的定量比較

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

3.3 消融實驗

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

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

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

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

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

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

3.4 運行時間

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

4 結論

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

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