結合全局語義優化的對抗性灰度圖像彩色化

萬園園， 王雨青， 張曉寧， 李荅群， 陳小林*

(1. 中國科學院 長春光學精密機械與物理研究所，吉林 長春 130033；2. 中國科學院大學，北京 100049)

1 引 言

圖像彩色化實質就是為灰度圖像的每一個像素分配顏色[1]。該技術在機器視覺、圖像修復和影視制作等方面具有重要的應用價值。現有的彩色化方法有傳統彩色化方法和基于深度學習的彩色化方法。

傳統的彩色化方法可分成兩種：基于局部顏色擴展的方法[1-4]和基于參考圖的顏色傳遞的方法[1-4]。其中，基于局部顏色擴展的方法需要人為在灰度圖像上進行顏色涂鴉，需要的用戶操作量較大，在顏色擴散過程中易出現邊緣誤擴散等問題。基于參考圖的顏色傳遞的方法需要選擇一張與待彩色化灰度圖相似的彩色圖進行顏色遷移，雖然此類方法減小了用戶操作量，但是對參考圖的依賴程度高。由于傳統算法受人工干預影響較大，且效率不高，在這個圖像海量化的信息時代已經不再適用。

目前，基于深度學習的彩色化算法已經能實現自動彩色化。如Cheng[1]等人利用卷積神經網絡來提取灰度圖特征，使用聯合雙邊濾波進行后續處理，實現了較好的彩色化效果。Iizuka[2]等人利用雙流結構，融合灰度圖像的局部特征和全局特征，并添加標簽分類來實現顏色預測，但生成圖像的顏色飽和度較低且細節損失較多。Zhang[3]等人將圖像彩色化問題轉化成對像素點顏色的分類問題，將飽和度高的顏色賦予高權重來獲得鮮艷的顏色，但是容易造成顏色溢出的問題。Larsson[4]等人使用VGG網絡來解釋場景的語義組成和對象的定位，但仍存在細節損失等問題。

近些年，生成對抗網絡(Generative Adversarial Network，GAN)因其特殊的網絡結構和訓練機制，在超分辨重建[5]、圖像合成[6]與風格轉換[7]等領域都取得了巨大的成功。受此啟發，GAN網絡也被用來解決圖像彩色化的問題。Isola等[8]人提出使用U-Net與GAN相結合的方法，在一定程度上提升了彩色化效果。后來Nazeri等[9]利用條件生成對抗網絡[8](conditional Generative Adversarial Network, cGAN)來還原彩色化的過程。Cao等[10]也通過在cGAN的多個卷積層中添加噪聲來增加彩色化效果的多樣性，但彩色化后顏色溢出的現象仍未改善。ChromaGAN[11]是Patricia等人提出來的一種最新的對抗性彩色化算法，他們把圖像場景分類引入GAN網絡，通過對顏色信息和類別分布的感知來實現灰度圖像彩色化，可改善場景彩色化錯誤的現象。由于以上基于GAN的彩色化方法均利用相對較淺的低維特征，對圖像的全局語義信息(例如場景的布局)理解不足，從而易導致顏色溢出等現象，且在特征提取時存在一定程度的細節損失。

針對以上問題，本文提出一種結合全局語義優化的生成對抗彩色化算法。通過對輸入圖像進行下采樣獲取多層特征和全局特征，并在跳躍連接中將得到的全局特征分別與多尺度層級特征融合來提升對圖像整體語義信息的理解能力，且在上采樣過程中引入通道注意力機制來抑制噪聲和降低冗余特征的權重，增強有用特征的學習能力[13]。同時搭建判別網絡，動態化地評價生成圖像的質量，進一步提升顏色的豐富程度。本文的損失函數在傳統顏色損失的基礎上引入帶梯度懲罰的Wasserstein生成對抗網絡[13-14](Wasserstein Generative Adversarial Net-work with Gradient Penalty, WGAN-GP) 的優化思想，用Wasserstein距離[15-16]度量真假樣本，以解決傳統GAN在訓練過程中出現的梯度消失和模式崩潰等問題。實驗證明，相比于目前主流算法，本文算法的彩色化效果有一定提升。

2 全局語義優化的對抗性彩色化算法

由于傳統的RGB模型并不能直觀地反映亮度信息，因此先把圖像轉換成Lab模型。生成網絡的輸入為亮度L通道。生成網絡主要包括3部分：下采樣、特征融合和上采樣。首先下采樣提取多尺度層級特征和全局特征，然后在跳躍連接中將得到的全局特征與每個尺度的層級特征融合，接著在上采樣過程中重構圖像，最終預測輸出三通道Lab的彩色圖。之后將生成的彩色圖和真實的彩色圖輸入判別網絡中進行識別。生成網絡和判別網絡在動態博弈中不斷增強性能，最終達到生成高質量圖像的目的。

2.1 生成網絡

本文的生成網絡在U-Net[18]網絡的基礎上進行改進，加深了網絡深度并引入自適應特征融合模塊和注意力機制。具體結構如圖1所示。

其中，為了減少最大池化帶來的信息丟失，本文采用卷積核為4，步幅為2，填充為1的卷積層逐步下采樣，依次獲得不同尺度的層級特征和全局特征。本文在跳躍連接添加特征融合模塊分別將全局特征和多尺度層級特征進行自適應融合。由于層級特征提取的是感受野相對較小的局部信息，而全局特征具有從整個圖像的感受野中提取的高維信息(如場景布局和類型等),因此全局特征可作為一種先驗信息來增強各個尺度的層級特征[2,17]，而且這兩種特征的融合形成了對圖像信息的高效表達，有利于生成顏色自然且細節豐富的圖像。

與此同時，在上采樣中將得到的全局特征逐漸恢復高分辨率來構建圖像。由于轉置卷積易產生棋盤格噪聲，本文采用雙線性插值和4×4的卷積層組合的方式代替轉置卷積來擴大圖像尺寸。隨后，按照U-Net的思想，在每個尺度下將已增強的低維層級特征通過跳躍連接與上采樣的高維特征拼接融合，共享低維特征的空間信息和細節信息[15,17]。為了提升算法特征學習的性能，添加了通道注意力機制[18]對融合之后的特征進行通道加權，抑制噪聲同時更關注與任務相關的特征通道。

2.1.1 自適應特征融合模塊

由于提取的層級特征感受野較小，主要包含局部特征信息。一般來說，全局特征反應圖像的整體結構，例如場景布局和類型等[17]。局部特征更加注重圖像的細節信息，反應圖像區域內更詳細的變化。兩者融合可以有效獲取豐富的特征信息，并且能增強算法全局語義信息的理解能力。因此，本文在不同尺度上通過特征融合模塊動態的將全局特征和包含局部信息的層級特征進行融合，如圖2所示。

圖2 特征融合模塊結構Fig.2 Structure of feature fusion block

假設層級特征圖fl的尺寸為H×W×Cl,全局特征圖fg的尺寸為1×1×Cg。首先將全局特征fg經過1×1的卷積層進行調整，該卷積層可自適應地從全局特征中提取最有用的信息來增強通道數為Cl的層級特征。然后把卷積之后的全局特征復制H×W次，使得到的全局特征尺度與待融合的層級特征尺度一致。最后將全局特征與層級特征進行拼接融合，輸出的特征尺寸為H×W×Cl。融合公式如下：

(1)

其中W是一個Cl×2Cl的權重矩陣，而b是一個維度為Cl的向量，W和b都是網絡的可學習參數。

2.1.2 注意力模塊

U-Net結構的優勢就是將下采樣的多尺度特征通過跳躍連接在上采樣過程中進行整合，使得語義信息豐富的高維特征和包含細節信息的低維特征互補，增強算法的性能。但是融合之后的特征層也包含了噪聲響應和冗余特征信息，會對圖像彩色化造成干擾。因此，為了能高效地從特征模塊中獲取有用特征信息，引入了通道注意力機制使網絡自動學習各個特征通道的相互依賴關系，抑制噪聲并增強有用特征通道的權重。注意力機制模塊結構如圖3所示。

圖3 注意力模塊結構Fig.3 Structure of attention block

輸入和輸出的特征圖尺寸均為H×W×C，首先對輸入的特征圖進行全局平均池化，得到一個1×1×C的特征圖。從某種程度來說，該特征圖具有全局感受野，可看作是對各個通道的描述。全局平均池化公式如下：

(2)

其中，uc(i,j)表示輸入特征圖的第c個通道內位置為(i,j)的像素值，Z為全局平均池化后的特征圖。隨后將特征圖Z經過兩個級聯的卷積層和激活層來學習各個特征通道的相關性。第一個卷積層對特征通道生成尺度為1×1×C/16的特征圖，第二個卷積層再把特征圖擴展回1×1×C，經過sigmoid函數激活輸出得到每個特征通道的重要性。最后得到的特征權重通過乘法加權在輸入的特征上，從而達到抑制噪聲和增強與任務相關的特征通道權重的目的[18]，在一定程度上提升了算法的性能。

2.2 判別網絡

判別網絡的目標是盡可能地區分生成圖像和真實圖像，本文判別網絡由5層全卷積層構成，詳細結構如表1所示。

表1 判別網絡結構Tab.1 Structure of discriminator

判別網絡前4層采用的是步長為2的4×4卷積層，同時每次下采樣之后特征圖的通道數加倍，每個卷積層后采用Leaky Relu激活函數。判別網絡的最后一層將特征圖映射成單通道輸出。判別網絡的輸入是生成網絡生成的彩色圖和與之對應的真實彩色圖，經過判別網絡處理之后返回輸入圖像是真或假的概率值。判別網絡對生成圖像進行評估判別，通過優化損失函數，反向傳播誤差調整網絡結構參數，以進一步改善生成網絡的彩色化效果。

2.3 目標損失函數的設計

傳統GAN大多基于JS散度進行優化，而JS散度易發散從而導致產生梯度消失的現象，不利于生成網絡進一步的學習。WGAN-GP網絡采用了Wasserstein距離[14]作為判別真假圖像的依據，且在損失函數中添加了梯度懲罰使得判別網絡滿足1-Lipschitz限制，以解決傳統GAN在訓練過程中產生的梯度消失和模式崩潰等問題。

本文將WGAN-GP的優化思想與顏色損失相結合來設計算法的損失函數。生成網絡的損失包括顏色損失和對抗損失兩部分，具體表達式如下：

Lg(G,D)=LC(G)+λ1Ladv(G,D)，

(3)

其中：LC表示顏色損失，Ladv表示對抗損失。由于Ladv∶LC=1 000∶1，通常情況下，當對抗損失為主導時，會使得網絡權重被帶偏，因此需要設置λ1用于平衡兩類損失。考慮到算法的主要任務是學習顏色信息，實驗設置λ1為1e-4，使得顏色損失占據主導地位。

對抗損失的具體定義如式4：

(4)

Pg表示生成彩色圖的數據分布，顏色損失LC表達式如式(5)所示，它是計算真實彩色圖與生成彩色圖的L1距離得到的。相比于L2損失，L1損失能較少生成圖像細節模糊的現象。其中Pr表示真實彩色圖的數據分布。

(5)

其次，判別網絡的損失函數引入了Wasserstein距離[14]，公式表達如式(6)所示。

(6)

公式前兩項表示真實數據分布與生成數據分布之間的Wasserstein距離。權重λ取10。Lgp表示梯度懲罰項，公式定義為：

(7)

(8)

3 實驗結果與分析

3.1 實驗環境與設置

本實驗在Pytorch深度學習框架下實現，利用TITAN RTX加速訓練。選用Place365場景數據集對網絡進行端到端的訓練，該數據集包含365個不同場景的180多萬張圖片。訓練之前先將數據集里少量灰度圖和場景灰暗的圖片濾除以保證算法能更好地學習顏色信息，處理之后最終選取1 687 424張圖片作為訓練集，10 000張圖像為測試集。

首先將輸入圖片顏色空間從RGB轉化成Lab，將L通道作為生成網絡的輸入。由于網絡存在尺度為1×1×C的全局特征，需要固定輸入尺寸的大小，實驗選用256×256的輸入圖像。算法采用Adam算法優化損失函數，初始學習率lr=0.000 1，β1=0.5,β2= 0.999，權重衰減為0.005。每批次大小為32，共訓練10個周期，訓練期間每迭代1 000個批次進行一次驗證輸出，觀察算法的彩色化效果。

3.2 評價指標

為了評估算法的彩色化效果，本文分別從定性和定量兩個方面進行分析。定性分析主要從視覺效果上進行評價，定量分析采用峰值信噪比(Peak Signal to Noise Ratio,PSNR)、結構相似性(Structural Similarity,SSIM)和信息熵(Information Entropy，IE)作為評價指標。峰值信噪比是一種使用廣泛的客觀評價指標，其數值越大表示失真越小。結構相似性是衡量兩幅圖像相似度的指標，而信息熵表示的是圖像包含信息的豐富程度，信息熵越大，圖像信息越豐富，圖像質量越好。

3.3 不同算法的結果對比

為了證明本文算法的有效性，將本文算法與現有彩色化算法[2-4,11]進行比較。實驗結果如圖4所示。

圖4 不同彩色化算法結果對比Fig.4 Result comparison of different colorization algorithms

首先針對顏色溢出的問題，例如第1行中，由于缺乏對圖像全局語義的理解，文獻[2,4]算法傾向于把跑道賦予天空的藍色，文獻[3,11]算法稍好一些，但也將原本紅色的跑道混合了少量藍色，而本文算法較好地將跑道的顏色恢復出來。在第2行，文獻[3]算法生成了更接近真實圖像飽和度高的紅色，而鮮艷紅色也違和地出現在輪胎上，文獻[2,4,11]算法得到的顏色飽和度較低且顏色不連續，本文算法在保持顏色飽和度較高的同時，也沒有出現顏色溢出的現象。在第3行，文獻[3]生成了紫色和黃色混合的花朵，呈現不自然的彩色化效果，文獻[2]花瓣的顏色飽和度較低且分布不均勻，文獻[4]的顏色鮮艷但是沒有還原花蕊的顏色信息，文獻[11]突出花蕊的細節信息，而有部分花瓣混合少量綠色。相較之下，本文算法生成的花朵顏色鮮艷且突出花蕊細節信息，在一定程度上能夠減少顏色溢出。

針對細節損失的問題，在第4行，本文算法還原了泳池顏色這個細節信息，而且周圍的草地樹木也呈現較好的彩色化效果。其他算法均沒有給泳池賦予恰當的顏色，泳池周圍的草地等也整體呈現灰褐色。在第5行中，其他算法都忽略水果筐的顏色細節，且對水果的顏色恢復較差，而本文算法不僅還原了水果筐的顏色細節，而且圖像整體顏色也更接近真實圖像。

通過以上對比可知，本文算法在保持較飽和顏色信息的同時，一定程度上也能夠緩解顏色溢出和細節損失等問題。

同時，為了從客觀的角度證明本文算法的優越性。采用3.2介紹的3種評價指標，對多種彩色化算法結果進行定量評估，對比結果如表2所示。表中計算的結果是對于所有測試集圖像的指標平均值。可以觀察到，對于相同的測試圖片，本文算法具有更好的評價指標，證明本文算法生成的彩色圖相比于其他幾種算法彩色化效果更佳。

表2 不同算法結果對比Tab.2 Result comparison of different algorithms

以上各種算法處理256×256尺寸大小的圖像平均所需時間如表3所示。

表3 不同算法處理時間Tab.3 Different algorithms processing time

3.4 消融實驗

本實驗采用不同的網絡結構進行訓練，來驗證GAN結構和各模塊對算法性能的影響。(1)原始U-Net+GAN；(2)U-Net+GF+GAN：GAN的生成網絡在U-Net基礎上添加全局特征優化模塊；(3)本文算法：U-Net+GF+Attention+GAN，生成網絡繼續添加通道注意力模塊；(4)本文算法去除GAN：U-Net+GF+Attention，單獨訓練本文算法的生成網絡G。前3種生成網絡復雜度依次升高，可分別驗證全局特征優化和通道注意力的作用，最后只訓練本文算法的生成網絡G，來驗證GAN結構對彩色化效果的影響。

由表4可知，添加全局特征優化模塊后，相比于原始的U-Net結構，各項指標均有所提升，說明全局特征與多尺度層級特征的融合有利于生成質量更高的圖像。通道注意力模塊的添加進一步增強算法的性能。而只訓練本文算法的生成網絡，評價指標均有所下降，證明GAN結構對圖像彩色化有一定提升效果。具體效果如圖5所示。

表4 各模塊作用驗證Tab.4 Effect validation of each module

原始U-Net結構的生成網絡容易出現顏色溢出的現象，如圖5(a)把房子的屋頂和部分墻面都賦予了天空的藍色，且忽略了房子旁邊的指示牌的顏色信息。添加全局特征優化模塊也相當于引入了全局語義優化，如圖5(b)屋頂和墻面的藍色去除，圖像的邊界更加清晰，指示牌也還原了較淺的紅色。說明全局特征優化模塊一定程度上增強了算法對圖像的整體信息的理解，有利于減少顏色溢出的現象。接著再添加完通道注意力模塊后(圖5(c)),指示牌和門等細節的顏色信息都有少許的加深。由此可知，全局特征優化模塊與通道注意力模塊對算法性能都有不同程度的提升效果。而在本文算法的基礎上去除GAN結構訓練得到的效果圖(圖5(d))，雖然緩解了顏色溢出的問題，但是顏色的飽和度和圖像質量都不如本文算法的效果圖(圖5(e))。這說明GAN網絡有利于生成彩色化效果更好的圖像。

圖5 不同網絡結構的結果對比Fig.5 Result comparison of different network structures

4 結 論

本文提出的彩色化算法是基于一種對抗策略，將全局特征與多尺度層級特征融合來捕捉圖像的全局語義信息，并且引入通道注意力模塊，通過抑制噪聲和增強與任務相關的特征通道的權重，以提高算法的性能。同時本算法的損失函數還在WGAN-GP的優化基礎上添加顏色損失，便于網絡訓練。實驗結果表明，該算法在主觀視覺上和評價指標上都取得了較好的效果。其中在Place365測試集上PSNR和SSIM指標分別達到24.455 dB和0.943。而且通過消融實驗證明了全局特征與多尺度層級特征融合模塊、注意力模塊和GAN網絡結構對算法性能的提升都起到了實質性的作用。相比于以往算法，本文算法在全局語義信息理解、圖像細節保持和顏色飽和度方面有較好的優勢。