基于注意力的光照感知水下圖像復原網絡

董鑫宇，朱 偉，黃詩芮，王光第，王柯儼，*，李云松

（1.西安電子科技大學 綜合業務網理論及關鍵技術國家重點實驗室，陜西 西安 710071；2.南京萊斯電子設備有限公司，江蘇 南京 210023）

0 引言

受到水中懸浮顆粒物的影響，光線傳播發生吸收和散射等作用，水下光學圖像普遍存在顏色失真、紋理細節丟失、對比度低和圖像模糊等問題。同時，水下環境光線昏暗，僅依靠自然光源難以成像，需要增加人工光源輔助。而這種非單一、不均勻的人工光源進一步增加了水下圖像復原的難度。混合光源環境下的水下光學成像原理如圖1所示。人工光源從成像設備處出發，經過成像設備到物體之間的距離照到物體上進行補光，這段傳輸距離中，人工光源同樣因為水下的吸收和散射作用被衰減。對物體進行光補償后這部分光同物體本身反射光一起，再次在水中傳輸過物體到成像設備之間的距離后被成像設備捕獲，在這一段傳輸距離中，補償后的反射光同樣會因為水下的吸收和散射作用被衰減。此外，進入成像設備的光線還有因散射進入成像設備的部分環境光。由于透射率隨場景深度指數衰減，則深度越大，透射率越小，即退化越嚴重。這類退化圖像影響了后續水下目標檢測等應用的性能。為了改善退化圖像的視覺質量，研究人員提出了諸多水下圖像清晰化方法。然而現有方法大多基于自然光源場景，很少考慮人工光源的影響。因此，研究包含自然光源和人工光源的混合光源環境下的水下圖像清晰化方法具有挑戰性和實用價值。

圖1 混合光源環境下的水下光學成像原理圖Fig.1 Principle of underwater optical imaging in hybrid light source environment

現有的考慮混合光源的水下圖像清晰化算法較少且基本是傳統算法。根據是否依賴成像模型，可分為傳統增強方法和傳統復原方法。傳統增強方法[1]通過融合色彩平衡和對比度增強，能夠有效增強不同光照條件下的水下圖像。但由于傳統增強方法并未考慮圖像的退化原理，忽略了退化程度與成像深度的關系，導致增強結果不夠自然真實，常常出現局部區域“過增強”或“欠增強”現象。傳統復原方法[2]通過最小信息損失先驗減弱人工光源的影響，借助顏色信息估計出透射率和環境光，代入水下成像模型中復原出清晰圖像。然而由于水下環境復雜多變，傳統復原方法難以得到高可靠的先驗條件，在一些場景下會因先驗失效而導致圖像復原結果不理想。

近年來，基于深度學習的方法在水下圖像處理領域已經取得了先進的性能。這類方法利用深度神經網絡的特征挖掘能力和非線性擬合能力，構建大量數據訓練網絡，從而學習退化圖像和清晰圖像的非線性映射關系。雖然部分深度學習方法考慮到人工光源的影響，但是受限于數據集和網絡結構，導致這些方法泛化性不足，不能很好地推廣到真實的混合光源水下場景中。

針對上述方法存在的問題，筆者提出了光照感知注意力編解碼器網絡用于混合光源下的水下圖像復原。首先，網絡引入光照感知作為先驗嵌入到網絡結構中，感知光源并均衡圖像的對比度。網絡本身基于多尺度的編解碼器結構，能夠有效集成多尺度特征。同時，注意力機制的引入使得網絡更關注需要復原的區域。筆者在具有豐富水下退化環境的數據集EUVP[3]上訓練和測試所提網絡，實驗結果表明，與現有典型方法相比，本文網絡在包含混合光源的水下圖像處理任務中能取得更為理想的效果。

1 相關工作

現有的水下圖像復原方法主要分為傳統方法和基于深度學習的方法。傳統的水下圖像復原方法主要是依據水下成像模型進行設計。但是，隨著計算機硬件成本的下降、計算能力的提升以及訓練數據集的豐富，基于深度學習的復原算法研究得到迅猛的發展，成為近年水下降質圖像清晰化研究的主流方向。

近年來，學者們提出許多基于深度學習的水下降質圖像處理算法，這些算法在圖像處理任務中得到的效果相對于傳統方法很好。基于深度學習的水下圖像處理算法方法分 3類：基于生成對抗模型（Generative Adversarial Networks，GAN）[4]的方法、基于水下成像模型的復原方法以及端到端的圖像復原算法。基于GAN的水下圖像復原方法訓練難度大，難以平衡生成器和鑒別器的性能。基于水下成像模型的深度學習算法通過CNN網絡估計透射率和環境光，然后利用成像模型完成圖像復原。這類方法一方面存在誤差累積問題，如微小的環境光估計誤差會導致復原結果出現明顯的色偏；另一方面受限于簡化的水下成像模型，不適用于混合光源場景。相比而言，端到端的復原方法不受水下成像模型的顯式約束，通過構建合適的網絡結構和損失函數，直接預測輸出復原后的水下圖像，其設計訓練更容易，復原效果更佳，適用范圍更廣。

由于端到端的復原算法優勢明顯，研究者們提出了很多代表性的工作，比如 MA等人提出的UIE-WD[5]，ANKITA等人提出的 Shallow-UWNet方法[6]和LIU等人提出的LANet[7]。UIE-WD利用離散小波變換將輸入圖像分解為多個子帶圖像，在頻率域借助豐富的高頻信息增強子帶圖像的細節和結構。該算法包含2個子網絡：多色空間融合網絡和細節增強網絡。其中，多色空間融合網絡將不同顏色空間的特征表示作為輸入，輸出顏色校正后的特征表示；細節增強網絡通過改進高頻子帶的圖像細節來解決原始水下圖像的模糊問題。Shallow-UWNet提出了一種輕量卷積神經網絡結構。該方法將原始輸入圖像通過跳躍連接與每個殘差塊的輸出進行拼接；跳躍連接對與原始輸入圖像相關聯的通道賦予更大的權重，確保每個殘差塊都能學習到原始圖像的基本特征。然而 Shallow-UWNet和UIE-WD并未考慮混合光源的影響，對混合光源場景水下圖像處理性能不佳，泛化性不足。考慮到水下環境中多樣的照明條件，LANet提出了一種基于注意力機制和自適應學習的水下圖像增強網絡。其中，并行注意力模塊用于關注光照特征和顏色信息；自適應學習模塊保留了淺層信息，自適應地學習重要的特征信息。

相較于 LANet借助注意力機制隱式地感知光照特征，筆者所提網絡直接引入光照圖作為約束加權到圖像特征，結合了傳統圖像先驗和卷積神經網絡的優點，解釋性更強，復原效果更好。

2 光照感知注意力編解碼器網絡

2.1 網絡結構

筆者提出的光照感知注意力編解碼器網絡（Illumination-aware Encoder-Decoder Net，IEDN，其整體架構如圖2所示。網絡將退化水下圖像I和對應的光照感知圖L作為輸入，輸出復原后的清晰圖像J。

圖2 網絡整體結構圖Fig.2 Overall structure of IEDN

在混合光源水下場景中，水質和照明條件復雜多變，導致退化類型多樣。尤其是人工光源的引入，可能導致圖像光照不均勻，局部區域過亮。如果不移除人工光源的影響，圖像質量恢復時容易出現過增強等現象。為了提高網絡在混合光源場景下的泛化能力，筆者引入了光照感知圖作為注意力先驗嵌入到網絡結構中，以均衡復原結果的對比度，提升整體視覺質量。光照感知圖[8]的指導原則是：賦予亮度暗的區域較高權重而亮度高的區域較低權重。權重和編解碼器網絡的輸出相乘，從而實現對圖像對比度的調整。網絡的整體處理流程

式中：fθ(·)為編解碼器網絡；θ為權重和偏置等模型參數；concat(·)為通道拼接操作。將退化圖像和光照圖通過通道拼接后輸入編解碼器網絡中，解碼端的輸出先和光照感知圖相乘均衡對比度，再通過長連接與輸入退化圖像相加，最后輸出清晰圖像。

IEDN的主體結構采用精心設計的三尺度并行編解碼器結構，利用從粗到細的策略恢復清晰水下圖像。高分辨率尺度保持了精細的空間細節和結構特征，低分辨率尺度擴大了感受野，近一步增強了網絡的細節重建能力。核心模塊包括：三尺度特征提取塊（Three-scale Feature Extraction Block，TFEB）、壓縮激勵殘差模塊（Squeeze-Excitation Res2block）、殘差雙注意力模塊（Residual Double Attention Block，RDAB）和特征融合模塊（Feature Fusion Block，FFB）。

三尺度特征提取塊組合了 3種不同核大小的卷積，拓展了感受野，有利于編碼器端提取豐富的結構特征。如圖 3所示，輸入特征圖在通道維被分割成 2個子特征圖，子特征圖分別通過不同核大小的卷積并行處理，然后在通道維進行拼接。設置了3種不同的卷積組合，分別是3×3/5×5、5×5/7×7和 3×3/7×7。三尺度特征提取塊還引入局部殘差，以有效緩解梯度消失等問題。壓縮激勵殘差模塊[9]能在粒度級別表示多尺度特征，并結合通道注意力高效引導網絡將注意力集中到待復原的區域。考慮到SE-Res2block參數量小且特征表示能力強，因此所提網絡在編碼器和解碼器中都采用了該模塊。

圖3 三尺度特征提取塊和殘差雙注意力模塊的結構Fig.3 Detailed structures of TFEB and RDAB

如圖3所示，殘差雙注意力模塊相較于普通殘差塊，引入了雙注意力，即并聯的通道注意力和像素注意力[10]。雙注意力的特征處理方式讓網絡將更多的注意力集中在不同水下場景中的重要像素點和重要通道上，可以靈活處理各種信息，對一些嚴重退化但有豐富細節的水下圖像的恢復效果提升顯著。

下采樣由平均池化層和卷積層組成，上采樣由雙線性插值和卷積層組成，卷積層的引入緩解了上下采樣操作導致的空間細節丟失問題。編碼器的中間特征通過特征融合模塊與解碼器的中間特征融合，增強了上下文信息的融合。特征融合模塊[11]通過逐像素點卷積和注意力模塊融合局部和全局上下文特征，相較一般的加法和拼接操作，能夠更好地解決尺度不連續的問題。

由于結合了上述模塊和光照感知的優點，IEDN具備強大的細節和結構保持能力，能夠復原出更為清晰的圖像紋理和更加真實的圖像色彩。

2.2 損失函數

本文方法采用平滑L1損失函數[12]、色彩感知損失函數[13]和邊緣感知損失函數加權方式引導網絡的學習。

平滑L1損失函數在誤差小于1時采用均方誤差形式，而其余情況均采用標準差誤差形式，從而有效兼顧了2種損失函數的優點，其表達式為

式中，ec代表誤差，即網絡復原結果和真值圖的差值。

色彩感知損失函數基于色彩通道距離指標，可校正水下圖像的色偏，引導復原圖像呈現更加自然的顏色。損失函數

式中：(x,y)為像素點的位置；H和W分別為圖像的長和寬；rn為所提網絡輸出圖的紅色通道與相應標簽圖紅色通道的平均值；Δrn、Δgn、Δbn分別為所提網絡輸出圖的紅色、綠色、藍色通道與相應標簽圖紅色、綠色、藍色通道的差值。該損失函數計算的顏色距離較常規的三通道等權重顏色距離，更符合人眼的觀感，由該損失函數引導學習得到的圖像顏色更加自然。

邊緣感知損失函數使用在 ImageNet預訓練好的 VGG16[14]網絡提取中間特征圖，量化復原結果和清晰圖像之間的視覺差異，有助于恢復水下圖像的細節信息和保持清晰的邊界。計算方式為

結合式（2）-（4），總的網絡訓練損失函數Lt由平滑L1損失函數、色彩感知損失函數和邊緣感知損失函數加權結合而成，即

該損失函數綜合考慮了網絡收斂速度以及圖像色彩和紋理細節恢復，引導網絡在達到理想結果。其中，1λ、2λ和3λ為權重系數。經過實驗，設置為1λ=0.7，2λ=0.15，3λ=0.15。

3 網絡訓練

在本節中，筆者首先介紹實驗細節。然后通過對比實驗來評估IEDN與其它方法的性能差距。最后通過消融實驗來驗證網絡結構和光照感知圖先驗對網絡性能的影響。

3.1 實驗細節

實驗采用EUVP的子集Underwater Scenes對網絡進行訓練和測試。該數據集場景豐富，共有2 185對配對圖像（包含混合光源環境的水下圖像及其真值圖），隨機劃分了1 600對訓練集、400對驗證集和185對測試集。

網絡在 Pytorch平臺上訓練和測試，并采用Adam優化器進行更新。從訓練集中每次抽取 16幅圖像送入網絡訓練，初始學習率 0.000 2。水下圖像輸入網絡前，會被隨機裁剪成256×256的分辨率，然后經過隨機旋轉和加噪被輸入到網絡中。整個網絡在NVIDIA GeForce RTX 3090 GPU上訓練200輪，耗時約5 h。

本文使用有參考評價指標峰值信噪比（PSNR）、結構相似度（SSIM）和色差公式CIEDE2000[15]對圖像質量進行評價。其中，PSNR和SSIM的指標越大越好，證明圖像質量越高；CIEDE的指標越小，說明生成的復原圖像越符合人的主觀感知。下面通過對比實驗和消融實驗來驗證所提算法的有效性。

3.2 對比實驗

筆者選取了5種算法作為對比方法，包括1種傳統圖像增強算法[1]，記為Fusion；1種傳統圖像復原算法[2]，記為 Li；3種自然光源下的深度學習算法，分別是2022年MA等人提出的UIE-WD[5]、2021年 ANKITA等人提出的 Shallow-UWNet方法[6]、2022年LIU等人提出的LANet[7]，分別記為UIE、Shallow和 LANet。為保證公平性，所有深度學習算法在本文使用的數據集上重新訓練。表1展示了EUVP測試集上，不同算法的有參考指標計算結果。

表1 對比實驗結果Table 1 Comparison of experimental results

實驗結果表明，本文所提算法在3個指標上的結果顯著優于其他算法，客觀上證明了我們算法的優秀性能。深度學習算法由于具有強大的學習能力，在性能上明顯優于傳統算法。UIE方法使用轉置卷積進行上采，而轉置卷積的不均勻重疊會導致復原結果存在棋盤效應，影響主觀效果和客觀指標。Shallow方法的優點是網絡簡單，但是同時導致網絡學習能力弱，所以性能比本文方法和LANet要差。LANet方法引入了雙注意力機制，性能優于Shallow方法，但缺少強力的多尺度特征提取能力，因此性能低于本文方法。由于本文方法使用了具有多尺度特征提取能力的TFEB和RDAB模塊，并且加入了光照感知圖先驗，在性能上明顯優于對比方法。因為TFEB和RDAB能夠高效提取豐富的結構細節特征，且光照感知圖能作為先驗約束網絡復原結果的對比度。

圖 4給出了混合光照環境下的水下圖像經過不同算法處理后的主觀效果，其中第7行為明顯的混合光源水下圖像。圖中偶數行圖像是奇數行圖像紅框內的局部放大圖像。從整體上來看，傳統方法的處理效果不理想，與深度學習方法有較大的差距。LI方法由于先驗估計不準確，復原結果的主觀質量不理想，引入了不自然的偽彩，如圖4第6行測試結果所示。而且，LI方法的結果在局部區域會出現明顯的過增強。Fusion方法雖然能夠有效提升圖像對比度，但是由于其未考慮圖像退化原理，部分增強結果不自然，如圖4中第2行測試結果所示。UIE方法由于棋盤效應，會使還原后的圖像出現不均勻重疊，主觀看起來就像是一塊塊的圖像拼接而成，如圖4第8行。Shallow方法由于網絡簡單，得到的紋理細節模糊，且增強效果有限，如圖4第2行。LANet方法由于對于光線變化敏感，會出現圖4第8行的噪點。總而言之，本文方法相對對比方法來說細節更豐富、色彩更符合人的主觀感覺；就光源的處理來看，本文方法復原的水下圖片光照更加自然，更接近標簽。

圖4 EUVP數據集上的復原結果對比Fig.4 Comparison of restoration results on EUVP benchmark

綜上，本文所提方法在混合光源場景下表現優異。

為了驗證本文方法與深度學習對比方法在算法復雜度上的差距，經過公平的實驗測試，在表2中列出了各項指標。

表2 算法復雜度Table 2 Algorithm complexity

由表2可知，UIE方法參數量最低，所用卷積和Activations最少，以性能換取計算復雜度。本文所提算法實現了性能和計算復雜度的均衡，所用FLOPs最少，在參數量和 Activations較低的情況下取得了最佳性能，推理速度也滿足實時性需求。

3.3 消融實驗

為了證明本文所提模塊和創新點的有效性，筆者對本網絡進行消融實驗，具體分為4個方案：不使用光照注意力、不使用壓縮激勵殘差模塊（Res2Block）、不使用RDAB模塊和不使用TFEB模塊的消融實驗，消融實驗結果如表3所示。

表3 消融實驗結果Table 3 Results of ablation experimen

由表3的實驗結果可知，相比無光照注意力的消融實驗，本文所提IEDN僅增加了極小的計算量實現了性能的提升。在4個模塊中，TFEB模塊對性能影響最大，無TEFB模塊的方案在PSNR指標上較IEDN相差0.9 dB，說明引入具有多尺度特征提取能力的TFEB模塊可以有效地提升網絡性能。無RES2模塊和無RDAB模塊的消融實驗與本文方法的指標差距也證明它們對網絡性能具有顯著的影響。從參數量上來看，TFEB模塊的影響最大，所需參數量接近本文方法的1/2。

4 結束語

本文針對現有水下圖像復原算法處理混合光源場景時的局限性，設計了一個具有多尺度特征提取能力的光照感知注意力編解碼器網絡。所提網絡具備良好的泛化性和魯棒性，能復原各種混合光源場景下的水下圖像，色調更自然，紋理細節更豐富，圖像質量更高，為后續水下高級視覺應用提供了有利條件。對比實驗的結果表明本文提出的光照感知注意力編解碼器網絡在主觀恢復效果和客觀指標評價上，均優于其他算法。消融實驗進一步證明了所提網絡的各個模塊和光照感知圖約束的有效性。雖然現在的網絡結構已經能夠取得比較理想的效果，但是仍然存在一些問題，如某些復原結果的邊緣特征不明顯。后續將通過增加邊緣輔助模塊，來增強網絡對邊緣特征的學習能力。此外還將探索一種有利于提升水下目標檢測算法的圖像復原方案。