基于UNDERWATER-CUT 模型的水下圖像增強算法

姚鵬,劉玉會

(中國海洋大學 工程學院,山東 青島,266100)

0 引言

視覺傳感器可以提供豐富的環境色彩信息,利用視覺信息對障礙物進行檢測和定位,進而對周圍環境進行三維重建,對水下機器人的導航及探索活動具有重要應用價值。然而,水下環境的圖像存在對比度低、清晰度不足等問題,對后續任務造成了阻礙。不同于空氣中的不清晰圖像(如霧化圖像),水中圖像主要是由水中折射導致的成像不清晰。光線波長影響光線在水中的衰減,波長越短,傳輸性能越好,因此波長短的光線占成像光線的大部分,導致了水下圖像大多呈藍綠色[1],難以直接對水下圖像進行目標識別等后續任務。同時水下存在大量漂浮顆粒物,導致環境光和背景光被吸收,進一步加劇了水下圖像對比度低和噪聲嚴重等問題,因此必須對水下圖像進行增強。

水下圖像增強通常采用基于物理模型和非物理模型2 種方式。基于物理模型的圖像增強方法對水下圖像的退化過程建立數學模型,并估算出模型的基本參數,然后通過該模型反演圖像的退化過程,獲得理想狀態下的未退化圖像[1-3]。然而,由于水下環境多樣,難以用具體模型建模,所以該類方法泛化性不強,無法修復多種水域圖像。

近年來,隨著深度學習的發展,基于深度學習的圖像增強算法得到了極大的進步。Li 等[4]提出了WaterGAN 兩階段模型,第1 階段構建生成式模型,根據成對的陸上圖像和深度信息,生成對應的水下圖像,第2 階段構建另一個生成模型,將生成的水下圖像和原本的陸上圖像作為成對的訓練集,訓練結束后,第2 階段的生成模型對水下圖像進行脫水增強。Li 等[5]提出了一種基于循環神經網絡的算法,該方法在循環生成對抗網絡(cycle generative adversarial network,CycleGAN)[6]的基礎上添加結構相似性(structural similarity,SSIM)損失約束[7],訓練時無需水下數據集和脫水數據集成對訓練,而且泛化性強,對于不同水域的水下圖像,均有不錯的增強效果。該方法通過維護水下圖像域X到脫水圖像域Y和脫水圖像域Y到水下圖像域X這2 個方向的映射,避免了多個X域圖像映射到同一Y域圖像的情形,但同時也引入了不必要的映射。為了修正顏色失真的水下圖像,僅需要維護從水下圖像到脫水圖像的單射,而不是2 個域圖像的雙射。李慶忠等[8]在CycleGAN 模型的基礎上添加了圖像強邊緣結構相似度約束,同時分兩階段訓練網絡,第1 階段進行改進的CycleGAN 模型的訓練,第2 階段,進行少量樣本的強監督訓練,得到了優于文獻的結果,但由于引入了強監督的樣本,不利于實際場景的應用。

針對文獻[5]存在不必要的雙射約束以及文獻[8]中存在的需要成對樣本參與訓練的缺陷,文中基于CUT 模型[9]提出了一種水下圖像顏色修正的算法(UNDERWATER-CUT 模型)。CUT 模型是CycleGAN 模型的一個增強模型,僅維護水下圖像域到脫水圖像域的單射,同時只需要非成對樣本參與訓練,在圖像域轉換等方面取得很好的效果,但是CUT 模型未針對水下圖像特點進行專門優化,對于水下圖像增強效果不佳。文中提出的UNDERWATER-CUT 算法為了保證顏色增強后圖像中物體的結構相似性,添加了SSIM 損失進行約束。此外,原始CUT 模型使用了InfoNCE 損失函數進行對比學習的約束[10],但是InfoNCE 函數形式復雜,不利于神經網絡收斂,因此文中使用alignment 和uniformity 2 個簡單函數的線性組合作為損失函數來進行模型的約束[9],得到了優于CUT模型的結果。

1 UNDERWATER-CUT 模型結構

圖1 為UNDERWATER-CUT 網絡結構,模型主體是對抗式生成網絡[11],由生成器和判別器兩部分組成,該網絡主要作用為通過非成對的水下數據集和脫水數據集(或陸上數據集)的無監督訓練,實現水下圖像域到脫水圖像域的轉換。不同于CycleGAN,文中的方法只需單向網絡即可保證高效轉換。設水下圖像域為X,脫水圖像域為Y,為了完成X域到Y域的轉換,需要同時訓練生成器G和判別器D。生成器需要是一個映射G:X→Y,一方面生成器構造映射y?=G(x),另一方面判別器衡量生成的質量D(y,y?),指引生成器梯度向更優的方向移動,當判別器不能分辨生成器產生的圖像時,此時達到最優。同時,為了防止模型將多個不同X域圖像映射為相同的Y域圖片,使用對比學習的方式約束圖像的生成,強迫模型學習到的映射G:X→Y為X域到Y域的單射。

圖1 UNDERWATER-CUT 模型網絡結構Fig.1 Network structure of UNDERWTAER-CUT model

生成器的具體結構如圖2 所示,總共分為3 部分子網絡:第1 部分由3 層卷積層網絡組成,主要作用為從 255×255×3的圖像中抽取出高維特征;第2 部分由6 層ResNet[12]組成,主要用來對前一階段提取出來的高維特征進行特征融合,學習到各層網絡之間的隱層關系;最后一部分也由3 層卷積神經網絡組成,其主要作用為將上游網絡層中的高維特征還原至3 個顏色通道的彩色圖像。

圖2 生成器模型網絡結構Fig.2 Network structure of generator model

為了同時捕捉到圖像中的顏色高頻分量和低頻分量,從而完整且客觀地權衡圖像的真實度,使用文獻[13]中的PatchGAN 判別器結構,將圖像分為 70×70的塊(patch),對于每一塊分別預測真假,最終將所有塊的預測值取均值得到最終對整幅圖像的預測值。這種判別器具體結構如圖3 所示,一共有5 層卷積神經網絡組成,和用于圖像二分類任務的常規模型不同,文中使用的PatchGAN 最后一層不使用全連接神經網絡,而是使用卷積層進行預測值的轉換,前4 層的主要作用是進行特征提取和融合,最后一層卷積神經網絡將隱層特征轉化為N×N×1的預測值,文中選取N=70。

圖3 判別器模型網絡結構Fig.3 Network structure of discriminator model

2 UNDERWATER-CUT 網絡損失函數和訓練策略

2.1 UNDERWATER-CUT 網絡損失函數

損失函數對于網絡模型至關重要,合適的損失函數意味著合適的約束,才能指引網絡學習到最優的情況,UNDERWATER-CUT 的損失函數主要由對抗生成損失、對比學習損失和SSIM 損失三部分組成。

1) 對抗生成損失

結合非成對數據集的水下圖像數據分布與脫水圖像數據分布,對抗生成損失的主要作用是使得生成器和判別器達成納什均衡,最優時生成器可以生成高質量的脫水水下圖像,即

2) 對比學習損失

對比學習損失的目的是最大化保留輸入的特征信息,將圖像分為多個塊,同時將生成器的前半部分Genc作為編碼器,便于壓縮圖像為一個向量表示,為了得到不同特征維度的信息,選取了L層來共同進行信息提取,為了最大化保留向量信息,會在每個l∈L層后面添加2 層的多層感知器(multilayer perceptron,MLP)網絡Hl,最終將L個編碼向量共同組合成為第l層的輸出向量,同樣的,對于生成器的輸出也存在原始CUT 使用InfoNCE函數作為基本損失函數,即

使用圖像塊作為基本樣本,原圖像對應位置作為正樣本,其他位置作為負樣本,同時對于多層進行計算。損失函數為

式(3)中對于每一層,將圖像向量分成Sl份,每一份圖像塊為第l層位置序號為s的生成圖像塊,為第l層位置序號為s的原圖像塊,為第l層位置不為s的其他生成圖像塊。式(2)由于函數形式復雜使得函數優化比較困難,難以收斂至最優值。使用文獻[14]提出的性質,InfoNCE 函數存在2 個性質,即相似樣本的相互拉近和其他樣本的均勻分布,故使用alignment 和uniformity 2 個簡單函數的線性組合來達到對比約束的效果,該損失函數為

將式(4)替換式(3)中的InfoNCE 函數,得出一個優化的PatchNCE 函數,稱為Opt-PatchNCE,即

3) SSIM 損失

為了更多的保留水下圖像的結構,使用結構一致函數,用于約束原水下圖像中的物體結構和增強圖像中物體結構一致,其定義為

式中:p為圖像塊的中心像素;x是X域中 大小為13×13 的圖像塊;y是G(x)域中大小為13×13 的圖像塊;ux,uy分別為x,y的均值;σx,σy分別為x,y的方差;σxy為x,y的協變量。SSIM 損失函數主要用于衡量x和G(x)的結構誤差,具體為

4) UNDERWATER-CUT 總損失

UNDERWATER-CUT 的總損失函數為對抗生成損失、對比學習損失和SSIM 損失三部分的線性組合,具體為

式中:λ1和λ2分別為約束對比學習損失和SSIM 損失的超參數,選擇λ1= 10,λ2=2.5。訓練時為了得到最優的生成器,按照下式進行總損失函數的優化

2.2 網絡模型訓練策略

每一次訓練迭代都分為2 個部分,第1 部分,對水下圖像域X和脫水圖像域Y分別隨機采樣m個樣本,固定生成器參數不變,對判別器的參數進行梯度下降的計算;第2 部分,只對水下圖像域X進行隨機采樣,由于第1 部分已經對判別器參數進行了優化,所以這一部分只優化生成器的模型參數。每次訓練判別器和生成器都能根據前一次迭代參數,得以進一步優化,這兩部分不斷互相指引對方的梯度向更優的方向優化,最終在有限次迭代次數中,二者達到最優,此時生成器增強的圖像最為清晰自然。

具體訓練配置如下:每次迭代選取8 個樣本作為1 批(batch),總的迭代次數(epoch)為200,網絡輸入為 255×255×3,梯度下降算法為Adam(adaptive moment estimation)算法,其中γ= 0.0002,β1= 0.5,β2=0.999。訓練集選取水下圖像增強基準(underwater image enhancement benchmark,UIEB)數據集[15],訓練時選取800 幅水下圖像和脫水圖像進行非成對訓練。文中算法具體實現使用PyTorch[16]框架,使用NVIDIA GeForce RTX 2080Ti 型號GPU 進行訓練,操作系統為Ubuntu 18.04。

3 實驗結果和分析

為了驗證方法的有效性,將從圖像生成效果和客觀指標評價兩部分來共同評價。圖像生成效果上將文中算法和原始CUT、CycleGAN 和SSIMCycleGAN 進行視覺效果上的比較,同時為了證明各個組成部分的有效性,進行了消融實驗,比較了SSIM-CUT 和UNDERWATER-CUT 的效果。在客觀指標評價上,由于單一指標評測有局限性,例如峰值信噪比(peak signal to noise ratio,PSNR)在某些情況下會出現圖像質量很高,但是PSNR 分數不高的情況[17],故采用了FID[18],水下圖像質量評價指標(undewater color image quality evaluation,UCIQE)[19]和PSNR 3 種指標來綜合評價,使用UIEB 數據集[15]作為訓練數據和測試數據,其中800 幅圖片為訓練集,90 幅圖片為測試集,訓練集和測試集中的圖片分辨率均為256×256。

3.1 圖像增強有效性分析

為了客觀衡量圖像增量后的效果,使用FID,UCIQE 和PSNR 3 種評價指標進行定量比較。FID主要被用來衡量生成模型生成的圖像增強結果,FID 的值代表了增強后的圖像集和陸上圖像集的距離,值越小代表距離越近,同時也代表了脫水后的效果越好。FID 將生成器生成的增強圖像和原輸入圖像看作2 個多維高斯分布,分別用N(μ,σ)和N(μw,σw)表示,然后計算2 個高斯分布的Wasserstein距離,具體為

式中,tr代表矩陣的跡。

UCIQE 指標是色彩濃度、飽和度和對比度的線性組合,用來定量評價水下圖像非均勻的色偏、模糊和低對比度的情況,UCIQE 的值和圖像質量成正比關系,值越高圖像質量越高,具體為

PSNR 反映了兩幅圖像之間的相似性,值越大表示增強后的圖像與陸上圖像越接近,即真實性越好,其定義為

式中,MSE為均方誤差函數,其具體定義為

式中:I為原始水下圖像;K為增強后的水下圖像;m和n分別為圖像的長度和寬度;MAXI為圖像點顏色的最大值,文中圖像每個采樣點用8 位表示,所以MAXI=255。

各方法增強結果見表1。根據表1 中結果可知,UNDERWATER-CUT 的FID 指標在5 種圖像增強算法中最低,說明增強后的圖像概率分布和原本水下圖像的概率分布最為接近,也就是生成的圖像保存信息最為完整;同時,文中方法的UCIQE分數也是最高的,說明在圖像的色度、飽和度和清晰度相比其他算法有著更好的表現;此外,文中方法的PSNR 分數接近SSIM-CUT 模型的分數,并都超過了其他算法,綜上,文中算法在PSNR 分數接近的情形下,FID 和UCIQE 均領先于其余算法,故可以認為經過UNDERWATER-CUT 增強的圖像質量優于其余算法。

表1 各方法增強后圖像指標對比Table 1 Comparison of the image indicators enhanced by each method

3.2 圖像增強效果對比

為了直觀比較各算法的增強效果,將原始圖像和各個算法增強后的圖像進行對比,結果如圖4 所示。原始CycleGAN 算法可以一定程度上增強水下圖像,但原始CycleGAN 算法未針對水下環境進行特殊處理,導致水下圖像中藍綠色失衡的問題仍然沒有被解決;添加SSIM 的CycleGAN 可以明顯提升圖像質量,但是由于使用了SSIM 損失函數,只約束了圖像中的物體結構,并未對圖像的色度和對比度進行處理,圖像顏色依然不夠真實自然;經過原始CUT 算法處理后,圖像中的物體邊緣發生了變化,和原始圖像相差較大,同時藍綠色失衡的問題依然存在;對CUT 添加SSIM 約束之后,可以看到增強后的圖像和原始圖像的物體結構一致,圖像較為清晰自然,但是色度和對比度偶有失衡,在結果中可以看到部分增強后的圖像背景被錯誤增強成了紅色;最后,UNDERWATER-CUT 效果最為清晰自然,顏色失真得到極大修正,因為使用了SSIM 損失函數又保證了圖像中的物體結構不改變,同時使用了相對比較容易優化的損失函數,模型更容易收斂到最優。

圖4 不同方法在相同數據集下的增強效果Fig.4 Enhanced results of different methods in the same dataset

4 結束語

文中提出了一種基于UNDERWATER-CUT模型的弱監督水下圖像增強算法,不需要水下圖像和脫水圖像的成對訓練,非常適合實際的水下任務。針對水下圖像特征,使用SSIM 約束確保修復后的圖像不丟失物體結構信息,保證修復后的圖像和修復前的結構信息的一致性,同時優化了對比學習損失,使得網絡進行梯度下降學習時更容易收斂到最優。實驗結果證明,經過UNDERWATERCUT 修復后的圖像在校正水下圖像顏色效果失真的基礎上,最大限度保存了水下物體的結構特征,增強后的圖像自然且清晰,有助于其他水下視覺任務如目標識別精度的提升。