改進U-Net網(wǎng)絡的水下圖像增強

李微，畢曉君

1. 哈爾濱工程大學 信息與通信工程學院，黑龍江 哈爾濱 150001

2. 中央民族大學 信息工程學院，北京 100081

水下圖像的質(zhì)量對海洋軍事、海洋資源開發(fā)等領域的研究具有至關重要的影響。然而，由于光在水中的衰減，造成獲取的水下圖像呈現(xiàn)出色偏、對比度低以及不清晰等問題，因此，對水下圖像進行增強成為研究者們關注的熱點。現(xiàn)有的水下圖像增強方法分為基于深度學習的水下圖像增強算法和基于傳統(tǒng)算法的水下圖像增強算法[1]。傳統(tǒng)的水下圖像增強方法能夠對特定場景下的水下圖像進行增強，但是傳統(tǒng)算法對于圖像中出現(xiàn)的不均勻光照、圖像霧化等不同場景下的圖像，容易出現(xiàn)局部區(qū)域過度增強而背景區(qū)域增強不夠且增強效率低的問題，實際應用價值有限。而近年來，深度學習方法[2]在對圖像進行處理上具有良好的效果以及較高的效率，因此，將深度學習方法應用到水下圖像增強中具有重要意義。

由于深度學習方法在提取圖像特征方面的表現(xiàn)良好，本文提出采用深度學習網(wǎng)絡作為水下圖像增強的網(wǎng)絡。首先本文提出采用生成對抗網(wǎng)絡合成水下圖像，制作數(shù)據(jù)集，之后在制作的數(shù)據(jù)集的基礎上訓練一個能夠充分提取水下圖像特征的改進U-Net網(wǎng)絡[3]用于真實場景下的水下圖像增強。

1 相關工作

深度學習方法需要規(guī)模龐大的數(shù)據(jù)集對模型進行訓練，然而這在深海環(huán)境中很難得到，因此需要制作水下數(shù)據(jù)集，接著利用深度學習網(wǎng)絡提取數(shù)據(jù)集中的水下圖像特征，使網(wǎng)絡能夠應用到真實場景下的水下圖像增強。

Wang等[4]提出一種端到端的基于卷積神經(jīng)網(wǎng)絡的水下圖像增強模型(underwater image enhancement-net，UIE-Net)，該方法首先利用水下成像模型合成了水下數(shù)據(jù)集，再將UIE-Net中的去霧子網(wǎng)絡和顏色校正子網(wǎng)絡分別用于圖像去霧和校正色偏2個任務，同時采用像素破壞策略來提取圖像局部模塊的固有特征，并且減少噪聲的干擾，從而極大地加快了模型的收斂速度和精度。Li等[5]提出一種能夠合成水下風格圖像的生成對抗網(wǎng)絡，稱為WaterGAN，WaterGAN通過無水圖像和深度圖像經(jīng)過真實水下圖像的渲染合成模擬的水下圖像，然后將模擬的水下圖像和無水真實圖像以及深度圖像都用于提供兩階段的深度學習網(wǎng)絡，網(wǎng)絡可以用來校正真實場景下的水下圖像。Fabbri等[6]提出一種魯棒性高的生成對抗網(wǎng)絡，該網(wǎng)絡首先使用循環(huán)一致性生成對抗網(wǎng)絡在未失真的圖像基礎上重建失真的圖像，然后對成對的水下圖像進行訓練，訓練出一種魯棒性高的網(wǎng)絡，該網(wǎng)絡可以將模糊的水下圖像轉換為清晰的高分辨率圖像。Li等[7]提出了一種基于水下場景先驗的卷積神經(jīng)網(wǎng)絡模型，稱為UWCNN。該模型結合水下成像模型和水下場景的光學特性，通過不同場景下的真實水下圖像合成了涵蓋不同水類型和退化程度的水下圖像退化數(shù)據(jù)集，利用相應的訓練數(shù)據(jù)對網(wǎng)絡進行訓練，得到的模型能用于真實水下圖像的增強。此外，基于該網(wǎng)絡的參數(shù)少、運算速度快，UWCNN模型可以擴展到水下視頻中進行逐幀增強。

現(xiàn)有的深度學習方法已經(jīng)在水下圖像增強領域取得了重大的突破，能夠對部分圖像進行有效地增強。但是現(xiàn)有方法合成的水下數(shù)據(jù)集與真實場景下的水下圖像還存在差距，以及現(xiàn)有深度學習模型不能充分提取不同風格水下圖像的全部特征，從而不能應用到不同風格的水下圖像中的問題，模型缺乏通用性。針對水下圖像合成質(zhì)量的問題，本文首先提出一個水下成像模型，然后使用生成對抗網(wǎng)絡模擬水下成像模型的過程，并對水下成像模型的參數(shù)進行自適應的學習，解決了基于水下成像模型人工設計參數(shù)的不足，同時訓練好的網(wǎng)絡恰好能擬合水下退化的真實過程。針對現(xiàn)有水下圖像增強方法不能充分提取水下圖像特征提取而導致模型通用性不足的問題，本文采用特征提取能力出色的U-Net網(wǎng)絡作為水下特征提取的基礎網(wǎng)絡，引入注意力機制[8]對水下圖像目標區(qū)域和背景區(qū)域的對比度進行增強，從而提升圖像整體質(zhì)量。

2 本文算法

2.1 水下數(shù)據(jù)集制作

目前，基于深度學習的水下圖像增強方法研究的難點為水下數(shù)據(jù)集的獲取。模型的訓練需要同一位置的無水圖像和退化圖像，但是無水圖像很難得到。因此，本文首先通過提出一個新的水下成像模型，利用生成對抗網(wǎng)絡，將無水圖像經(jīng)過真實水下圖像的渲染生成無水圖像對應的合成水下圖像，從而建立成對的水下數(shù)據(jù)集。

2.1.1 水下成像模型

光在水中傳播過程中，由于水體對光的吸收和散射造成光的衰減，水下環(huán)境中還存在水中懸浮顆粒、浮游微生物和無機鹽等都能吸收光能，同時產(chǎn)生光的散射效應而改變光的傳輸路徑，從而使在水下獲取的圖像出現(xiàn)色偏、不清晰以及對比度低等問題。經(jīng)典的水下成像模型為

式中：x為 圖像中的每個像素；J(x)為無水圖像；t(x)為傳輸效率；E(x)為退化圖像；η為衰減系數(shù)；d(x)為目標到水下相機的距離；B與圖像遠處的背景光相關。

經(jīng)典模型只考慮了光在水中的直接衰減和后向散射衰減的影響，然而，水下環(huán)境中由于溫度低，使得水中的小水滴和水汽凝結形成水霧，水霧會使得光傳輸衰減。當水下相機獲取圖像時，圖像場景深處的霧化程度和場景淺處的霧化程度不一致，從而造成水下圖像的不同地方因為霧的影響而造成的能見度低下的程度不一樣。而普通的水下模型沒有考慮場景深度與霧化程度的影響，在場景深度變化大的圖像中合成效果差，難以模擬出圖像在場景深處的霧化效果。因此，其合成的水下圖像就不能很好的模擬真實水下圖像成像的特性。本文結合場景深度對圖像霧化的影響，并對直接衰減和后向散射衰減提供2個不同的衰減系數(shù)，提出的水下成像模型為

式中：x為 圖像中的每個像素；J(x)為無水圖像；t(x) 為水中自然衰減的傳輸效率；β(λ)為水中不同波長的光在水下的衰減系數(shù)；d(x)為目標物到水下相機的距離；t′(x)為水中霧化情況下光的傳輸效率；α為衰減系數(shù)；I(x)為退化后的圖像。

2.1.2 基于生成對抗網(wǎng)絡的水下圖像生成模型

近年來生成對抗網(wǎng)絡在圖像生成領域表現(xiàn)良好，本文利用生成對抗網(wǎng)絡模擬水下成像模型的衰減、散射以及霧化進行生成對抗網(wǎng)絡結構的設計，能夠自適應地學習水下成像模型的參數(shù)，網(wǎng)絡結構如圖1所示。具體方法為采用無水深度圖像作為生成器的輸入，利用真實場景下的水下圖像作為判別器的監(jiān)督圖像，來引導生成器生成圖像的風格，從而使得生成器生成的圖像與真實場景的水下圖像難以區(qū)分，達到生成的圖像與真實的圖像具有相同水下風格的目的，從而制作成對的水下數(shù)據(jù)集。合成的2種風格的圖像如圖2所示，圖2(a)列為無水圖像，圖2(b)列監(jiān)督圖像為淺水圖像對應的合成水下圖像，圖2(c)列監(jiān)督圖像為深水圖像對應的合成水下圖像。

圖1 水下圖像生成網(wǎng)絡整體框架

圖2 2種風格的合成圖像示例

2.2 水下圖像增強模型

本文提出在U-Net網(wǎng)絡的基礎框架上引入注意力機制作為水下圖像增強的模型，模型整體框架如圖3所示。

圖3 水下圖像增強整體框架

2.2.1 特征提取

本文水下圖像增強模型一共4層，模型輸入像素為256×256的彩色圖像，壓縮路徑在第1～4層的卷積核數(shù)量依次為32、64、128、256；2層之間進行1個2倍收縮的下采樣，每1層圖像大小依次為256×256、128×128、64×64、32×32；擴增路徑與壓縮路徑對稱，在2層之間采用1個2倍擴增的上采樣，在同層之間使用跳躍連接融合擴增路徑與收縮路徑的特征信息，在第3層的跳躍連接之間引入1個注意力機制模塊，提高深層抽象特征的學習能力，注意力機制模塊的輸入為收縮路徑第3層的輸出和經(jīng)過上采樣的第4層輸出，獲得抽象特征的注意力概率分布之后，注意力機制的輸出再與經(jīng)過上采樣的第4層輸出進行信息融合。收縮路徑與擴增卷積的卷積層的卷積核大小均為3×3，每個卷積層后面跟1個非線性激活函數(shù)Relu，模型輸出的卷積核數(shù)量為3。

2.2.2 損失函數(shù)

在水下圖像增強網(wǎng)絡中，選擇一個合理的損失函數(shù)引導網(wǎng)絡增強后的圖像能夠逼近監(jiān)督圖像尤為關鍵，同時可減少網(wǎng)絡的訓練時間，幫助網(wǎng)絡更好地訓練。L1損失廣泛應用于圖像增強領域中，同時圖像的結構相似程度也是衡量圖像增強質(zhì)量的一個重要指標，因此，采用式(3)作為總的損失函數(shù):

式中：Lloss為總損失；L1為絕對誤差損失；Lssim為 結構相似度損失；T(x)表示監(jiān)督圖像中像素x的值；G(x)表示增強圖像中像素x的值；?為2個損失各自的權重系數(shù)。

3 實驗結果與分析

3.1 實驗環(huán)境與參數(shù)設置

本實驗的硬件環(huán)境配置Intel?CoreTMi7-7700 CPU@3.6 GHz×8，64位處理器，采用GeForce GTX 1080TiGPU進行運算加速，操作系統(tǒng)為64位Ubuntu 16.04，采用TensorFlow2.0框架以及python3.6進行程序實現(xiàn)。

水下圖像生成模型的訓練代數(shù)為10，批量大小為8，優(yōu)化器為Adam，指數(shù)衰減因子0.5，學習率2×10?4；水下圖像增強模型的訓練代數(shù)為200，批量大小為32，優(yōu)化器為Adam，指數(shù)衰減因子0.9，學習率2×10?4。

3.2 數(shù)據(jù)集選取

采用NYU Depth dataset V1[9]和NYU Depth dataset V2[10]中的無水圖像以及對應的深度圖像共3 728張，經(jīng)過水下機器人大賽官方提供的淺水區(qū)域和深水區(qū)域2種風格的真實水下圖像的渲染，合成無水圖像對應的水下圖像，從而構造了水下數(shù)據(jù)集。將合成水下數(shù)據(jù)集分2部分，用3 712對合成圖像進行水下圖像增強模型的訓練。測試集包括3個數(shù)據(jù)集：數(shù)據(jù)集1為16對合成數(shù)據(jù)集；數(shù)據(jù)集2為水下機器人官方提供的數(shù)據(jù)集，包括淺水區(qū)域圖像2 069張，深水區(qū)域圖像2 173張；數(shù)據(jù)集3為多場景的UIEBD數(shù)據(jù)集[11]，共890張。

3.3 測試結果與分析

為了驗證本章所提出模型的有效性和先進性，進行了一系列對比實驗，對比方法包括以HE[12]、Fusion[13]為代表的基于像素值重新分配的水下圖像增強方法，以UDCP[14]、IBLA[15]、MIP[16]為代表的基于水下成像模型的水下圖像增強方法，以UWCNN[7]為代表的基于深度學習的水下圖像增強算法。

3.3.1 算法有效性驗證

首先，為了驗證添加注意力機制模塊的有效性，使用數(shù)據(jù)集1中的訓練集分別對添加注意力機制的U-Net水下圖像增強網(wǎng)絡與未添加注意力機制的U-Net水下圖像增強網(wǎng)絡進行訓練。使用合成的水下圖像作為網(wǎng)絡的輸入，將相應的無水圖像作為網(wǎng)絡的監(jiān)督圖像，通過訓練集的訓練得到訓練好的模型。由于數(shù)據(jù)集1中的合成水下圖像具有對應的無水圖像作為監(jiān)督圖像，因此，在數(shù)據(jù)集1中的測試集中采用有參考圖像的評價標準峰值信噪比(peak signal-to-noise ratio，PSNR)和結構相似度(Structural SIMilarity，SSIM)進行不同方法測試效果的對比，結果如表1所示。

表1 不同方法的PSNR和SSIM結果對比

表1中，Ours1為未添加注意力機制的U-Net水下圖像增強網(wǎng)絡的實驗結果，Ours2為添加注意力機制的U-Net水下圖像增強網(wǎng)絡的實驗結果。對比二者可以看出，在U-Net網(wǎng)絡中添加注意力機制模塊能夠有效提升圖像的SSIM值與PSNR值。同時，由表1可知，基于深度學習的水下圖像增強方法通過對訓練集特征的提取，在與訓練集同分布的數(shù)據(jù)集1上的增強表現(xiàn)優(yōu)于其他方法，證明了深度學習方法在數(shù)據(jù)集1上的優(yōu)越性。

3.3.2 算法先進性驗證

為了驗證模型在不同風格水下圖像中的泛化能力，采用數(shù)據(jù)集2這種真實場景下的水下圖像作為測試集能夠更有效地說明模型對真實場景下圖像的增強效果。而本文所提出的模型是在數(shù)據(jù)集1中的水下風格合成圖像上進行訓練的，而數(shù)據(jù)集1的水下風格來源于數(shù)據(jù)集2，因此，模型經(jīng)過訓練能夠學習到數(shù)據(jù)集2中2種水下風格圖像的特征，以此作為先驗信息，從而在類似風格的水下圖像中增強效果明顯。為了驗證本文所提出的模型能夠對不同風格的水下退化圖像進行增強的通用性，在含有豐富水下場景的數(shù)據(jù)集3上進行實驗。如圖4—圖6所示。

圖4 不同算法對示例1的增強效果

圖5 不同算法對示例2的增強效果圖

圖6 不同算法對示例3的增強效果

圖4—圖6分別為不同算法在數(shù)據(jù)集2中淺水區(qū)域的水下圖像增強效果示例、數(shù)據(jù)集2中深水區(qū)域的水下圖像增強效果示例以及數(shù)據(jù)集3中水下圖像增強效果示例。由圖4—圖6可以看出，HE方法能夠顯著提升圖像的亮度和色彩度，但是對水下物體會出現(xiàn)由于過度增強而導致失真的情況；Fusion方法增強之后的圖像在場景深度深處增強不徹底，同樣會出現(xiàn)失真的情況；UDCP方法增強的圖像偏暗；IBLA方法增強的圖像會使顏色更深，MIP方法僅能夠降低水下衰減的程度，對圖像全局改善有限；UWCNN增強的圖像能夠貼近沒有水下衰減的情況，但是圖像整體偏暗；本文算法能夠顯著提高圖像的對比度與清晰度，校正圖像色偏，徹底去除水下衰減的影響，尤其在場景深度深的背景區(qū)域，增強效果明顯優(yōu)于其他方法。

為了客觀驗證本文算法的先進性，采用無參考圖像評價指標UIQM這種能夠對圖像色彩豐富度、圖像清晰度以及圖像對比度綜合考慮的評價指標對圖像質(zhì)量進行評價。圖4—圖6中圖像的UIQM值如表2所示。

表2 不同方法的UIQM值對比

由表2可以看出，本文算法在3張圖中的UIQM值都為最高，展現(xiàn)出本文算法的優(yōu)越性，同時，HE算法的UIQM值也較高。

對3個數(shù)據(jù)集的所有水下退化圖像進行測試，得到的UIQM均值如表3所示。

表3 不同方法在3個數(shù)據(jù)集上的UIQM均值對比

UIQM值近似看作與圖像質(zhì)量呈正相關。表3中，淺水數(shù)據(jù)集和深水數(shù)據(jù)集中的HE方法的UIQM均值最高，然而HE方法是通過像素的重新分布，對圖像的色彩度和亮度進行過量提高，從而UIQM值也偏高，同時圖像也會出現(xiàn)失真的現(xiàn)象，因此UIQM值不一定能夠準確地衡量圖像質(zhì)量。基于主觀視覺與客觀指標的綜合評價，本文算法能夠對3個數(shù)據(jù)集的水下退化圖像進行有效地增強，且使增強后的圖像符合人類視覺感知，模型的泛化能力強。

表4為不同算法在數(shù)據(jù)集3中890張圖像上生成增強圖像的平均速度。結果表明，基于深度學習的方法包括UWCNN與本文算法的增強效率優(yōu)于傳統(tǒng)算法。

表4 不同算法在單張圖像上的平均增強速度

結果表明，基于深度學習的方法包括UWCNN。

4 結論

本文首先提出一個新的水下成像模型用于合成具有水下特征的數(shù)據(jù)集，通過改進U-Net網(wǎng)絡對合成的數(shù)據(jù)集進行訓練得到增強網(wǎng)絡模型。通過在3個不同的數(shù)據(jù)集上進行算法的有效性和先進性驗證。結果表明，經(jīng)過本文算法增強后的圖像提高了圖像的對比度和清晰度，同時能夠有效校正圖像色偏，提高水下圖像的整體質(zhì)量，能夠有效地對不同場景下的水下圖像進行增強，模型通用性強。同時在主觀視覺上，增強后的圖像符合人類視覺感知，貼近真實。且增強的效率高，為水下圖像的實時增強提供了可能。