基于改進(jìn)U-Net網(wǎng)絡(luò)的圖像自動白平衡算法

全雪峰

（南陽醫(yī)學(xué)高等專科學(xué)校 衛(wèi)生管理系，河南 南陽 473061）

0 引言

顏色校正問題對于諸如顏色識別、圖像分類、圖像目標(biāo)檢測等應(yīng)用來說非常重要，顏色的偏差可能導(dǎo)致圖像處理結(jié)果錯誤。人眼能夠自適應(yīng)光源的變化，并還原場景真實(shí)顏色，而圖像傳感器不具備自適應(yīng)的功能。為了解決這個問題，數(shù)字?jǐn)z像機(jī)等使用白平衡方法對圖像顏色進(jìn)行修正，使其符合人眼視覺。

目前，顏色校正研究取得了一定進(jìn)展，經(jīng)典方法如灰度世界法[1]、完美反射法[2]、動態(tài)閾值法[3]以及這些算法的改進(jìn)[4-5]。這類算法原理比較簡單，容易實(shí)現(xiàn)，但是當(dāng)先驗(yàn)條件不滿足時算法會失效。例如當(dāng)含有大量的單色塊時，灰度世界法會失效；當(dāng)圖像中不含白色點(diǎn)時，完美反射法會失效。Cardei等人[6]提出利用BP神經(jīng)網(wǎng)絡(luò)進(jìn)行光照估計(jì)，進(jìn)而校正顏色。神經(jīng)網(wǎng)絡(luò)方法比較靈活，學(xué)習(xí)的目標(biāo)不同就可以使網(wǎng)絡(luò)擁有不同的功能，其缺點(diǎn)是輸出結(jié)果容易陷入局部極小[7]。

近年來基于卷積神經(jīng)網(wǎng)絡(luò)的深度學(xué)習(xí)算法得到快速發(fā)展[8-11]。如Shi等人[12]提出用于圖像光照估計(jì)DS-Net網(wǎng)絡(luò)；Hu等人[13]提出基于置信度加權(quán)池的全卷積顏色恒常性算法；Zhu等人[14]對用于圖像風(fēng)格遷移的網(wǎng)絡(luò)結(jié)構(gòu)進(jìn)行改進(jìn)，并用于圖像顏色校正；Shen等人[15]根據(jù)Retinex顏色恒常性理論提出MSR-net網(wǎng)絡(luò)，用于解決低光照下顏色增強(qiáng)問題；Simone等人[16]提出一種準(zhǔn)無監(jiān)督的深度卷積神經(jīng)網(wǎng)絡(luò)顏色恒常性算法。本文設(shè)計(jì)了一種新的端到端的圖像顏色校正網(wǎng)絡(luò)。實(shí)驗(yàn)結(jié)果表明，本文算法可以有效校正圖像色偏。

1 U-Net網(wǎng)絡(luò)介紹

U-Net網(wǎng)絡(luò)[17]是一種被廣泛應(yīng)用于醫(yī)學(xué)圖像分割的深度卷積神經(jīng)網(wǎng)絡(luò)，由收縮路徑和擴(kuò)張路徑兩部分組成，采用編碼器/解碼器結(jié)構(gòu)。在收縮路徑中包括4個階段，每階段由2個3×3卷積核和步長為2的2×2最大值池化構(gòu)成，使用ReLU激活函數(shù)，通過收縮路徑得到低分辨率的高維特征。擴(kuò)張路徑也包括4個階段，每個階段用2×2的卷積核進(jìn)行反卷積操作，然后用兩個3×3的卷積核進(jìn)行卷積運(yùn)算。為彌補(bǔ)收縮路徑中的信息丟失，在每階段的反卷積之后，與收縮路徑中對應(yīng)層級的特征圖相拼接。最后通過一個1×1卷積輸出結(jié)果。U-Net網(wǎng)絡(luò)可以使用很少的圖像進(jìn)行端到端的訓(xùn)練，減小了對樣本量需求的壓力。

2 基于改進(jìn)U-Net網(wǎng)絡(luò)的顏色校正方法

2.1 網(wǎng)絡(luò)結(jié)構(gòu)

本文對U-Net網(wǎng)絡(luò)進(jìn)行改進(jìn)，設(shè)計(jì)了一種新的用于顏色校正的端到端的網(wǎng)絡(luò)結(jié)構(gòu)，如圖1所示。

圖1 網(wǎng)絡(luò)結(jié)構(gòu)

由圖1可見，整個網(wǎng)絡(luò)主要有由三部分組成：（1）用于提取特征的編碼器；（2）用于提取圖像細(xì)節(jié)信息的稠密連接塊（Dense Block）；（3）用于恢復(fù)圖像的解碼器。整個網(wǎng)絡(luò)類似于U-Net網(wǎng)絡(luò)結(jié)構(gòu)，

編碼器：本文采用4層卷積構(gòu)成一個收縮路徑（左側(cè)），除第一層外，在其他3層卷積塊之前分別放置一個最大值池化層，以縮小特征圖尺寸，降低網(wǎng)絡(luò)訓(xùn)練參數(shù)，加快訓(xùn)練速度。每一層的卷積采用相同尺寸（3×3）的卷積核。

Dense Blocks：不同于經(jīng)典的U-Net網(wǎng)絡(luò)，本文在編碼器和解碼器之間加入被稱為Dense Blocks的卷積塊，該卷積塊來源于AOD-Net網(wǎng)絡(luò)[18]。AOD-Net是一種用于除霧的輕量型神經(jīng)網(wǎng)絡(luò)，它通過端到端的CNN直接生成清晰圖像，而且通過多尺度特征映射之間的合并，提高了模型的性能。本文引入AOD-Net網(wǎng)絡(luò)結(jié)構(gòu)，但去掉原網(wǎng)絡(luò)中最后的減法運(yùn)算。實(shí)驗(yàn)表明，通過加入Dense Blocks，可以提高恢復(fù)圖像的清晰度。

解碼器：在解碼器中，本文通過上采樣和具有相同尺寸（3×3）卷積核的卷積塊，構(gòu)成一個與收縮路徑相對應(yīng)的擴(kuò)展路徑（右側(cè) ）。

實(shí)驗(yàn)證明，經(jīng)過上述編碼器、Dense Blocks和解碼器后恢復(fù)的圖像清晰度較低，失去部分圖像輪廓信息。為此，本文采用U-Net網(wǎng)絡(luò)的思想，將來自收縮路徑的高分辨率特征與上采樣層輸出在通道維上相拼接，這樣有效地彌補(bǔ)了由于前面的池化操作帶來的信息損失。在網(wǎng)絡(luò)的最后一層，使用一個1×1的卷積來映射8個特征圖，輸出最終校正顏色后的圖像。

2.2 損失函數(shù)

本文在訓(xùn)練網(wǎng)絡(luò)時使用的損失函數(shù)是均方誤差（Mean Squared Error，MSE）[14]，定義如下：

式中，Io是網(wǎng)絡(luò)輸出圖像，Ir是參考圖像，P是圖像塊，N是像素總數(shù)。

3 實(shí)驗(yàn)

本文的實(shí)驗(yàn)平臺為Deepin Linux 11.5，CPU為intel i5 7600，GPU為Nivida GForce GTX 1050Ti，16G內(nèi)存，采用Pytorch深度學(xué)習(xí)框架實(shí)現(xiàn)。

3.1 數(shù)據(jù)集

本文使用文獻(xiàn)［19］提供的數(shù)據(jù)集Set1進(jìn)行網(wǎng)絡(luò)訓(xùn)練和測試，該數(shù)據(jù)集源于被廣泛用于光照估計(jì)的NUS數(shù)據(jù)集[20-21]和Gehler數(shù)據(jù)集[22]，包含了由不同型號的相機(jī)采集和采用不同白平衡設(shè)置后渲染的62535張sRGB圖像及對應(yīng)的原始圖像。本文從Set1數(shù)據(jù)集中隨機(jī)選擇56669張圖像訓(xùn)練網(wǎng)絡(luò)，剩下的5866張圖像用于測試。在測試階段，本文還使用了文獻(xiàn)［19］提供的Set2數(shù)據(jù)集和Cube+數(shù)據(jù)集，前者包含了2881張渲染后的圖像和對應(yīng)的原始圖像，后者包含了10242張渲染圖像和對應(yīng)的原始圖像。由于本文網(wǎng)絡(luò)是全卷積網(wǎng)絡(luò)，模型可以接受任意尺寸的圖像，但為保證處理速度，在將圖像輸入網(wǎng)絡(luò)之前將其尺寸調(diào)整為256×256；測試時，將輸出的圖像恢復(fù)到其原始尺寸。

3.2 訓(xùn)練參數(shù)

考慮到將學(xué)習(xí)速率設(shè)置過大，難以找到最優(yōu)解，設(shè)置過小，網(wǎng)絡(luò)收斂緩慢，本文設(shè)置初始學(xué)習(xí)率為0.01，權(quán)值衰減為0.0001，batch=8，采用Adam優(yōu)化器。由實(shí)驗(yàn)發(fā)現(xiàn)，經(jīng)過10個epoch后就可以充分訓(xùn)練網(wǎng)絡(luò)，如圖2所示。

圖2 Epoch-loss曲線

4 結(jié)果分析

4.1 評價指標(biāo)

本文采用均方誤差MSE、平均角度誤差（mean angular error, MAE）[19]和ΔE 2000[23]作為評價指標(biāo)。

4.2 不同卷積設(shè)計(jì)對顏色校正效果的影響

本文針對編碼器和解碼器的卷積核大小、層數(shù)進(jìn)行了多種設(shè)計(jì)，具體為{1, 3, 5, 7, 2}、{3, 2}、{3, 1}，其中{1,3, 5, 7, 2}代表編碼器的卷積核大小從上到下依次為1、3、5、7（對應(yīng)的解碼器從下到上依次為1、3、5、7），每層有2個卷積；{3, 2}代表編/解碼器的卷積核大小統(tǒng)一設(shè)為3，每層有2個卷積；{3, 1}代表編/解碼器的卷積核大小統(tǒng)一設(shè)為3，每層1個卷積。圖3所示為選擇不同參數(shù)值后的顏色校正結(jié)果，表1所示為對應(yīng)的評價指標(biāo)值。

由圖3和表1可知，當(dāng)網(wǎng)絡(luò)設(shè)計(jì)為{3, 2}時，顏色校正后的圖像質(zhì)量最好，原因是網(wǎng)絡(luò)采集到的圖像特征最多。在隨后的實(shí)驗(yàn)中，采用的都是編碼器和解碼器的所有層的卷積核大小為3，每層有2個卷積。

圖3 不同參數(shù)圖像校正結(jié)果

表1 不同參數(shù)圖像校正結(jié)果

4.3 不同算法的視覺比較

圖4所示為部分色偏圖像用灰度世界法[1]、Quasi-U CC算法[16]和本文算法校正顏色后的視覺比較。由圖4可以看出，除了圖片f以外，灰度世界法的顏色校正效果均不理想，尤其是對圖片d和圖片e的校正結(jié)果出現(xiàn)明顯的色偏，這是因?yàn)樵撍惴僭O(shè)在標(biāo)準(zhǔn)光源下的場景，其平均反射是無色差的，這樣當(dāng)場景包含有大片單色區(qū)域時，校正的結(jié)果會出現(xiàn)色偏。Quasi-U CC算法的顏色校正效果優(yōu)于灰度世界法，但對圖片d，基本未校正顏色，對于圖片g，校正后的圖片較暗。本文算法的顏色校正結(jié)果均取得了較好的視覺效果，優(yōu)于灰度世界法和Quasi-U CC算法。

圖4 不同算法顏色校正效果比較

4.4 不同算法客觀比較

為了從客觀上驗(yàn)證本文算法的有效性，將本文算法與現(xiàn)有顏色校正算法（灰度世界法[1]、完美反射法[2]、FC4算法[13]、Quasi-U CC算法[16]）進(jìn)行了對比。表2所示為幾種不同算法的評價指標(biāo)值，這些值是從文獻(xiàn)中獲取的，或者是通過執(zhí)行公開可用的實(shí)現(xiàn)獲得的。

由表2可以看出，本文方法在不僅在Set1數(shù)據(jù)集上取得了較好的成績，而且在不同數(shù)據(jù)集上也取得了較好的成績，并且在很大程度上優(yōu)于現(xiàn)有算法，同時顯示出更穩(wěn)定的性能。由表2還可以看出，基于神經(jīng)網(wǎng)絡(luò)的算法（FC4、Quasi-U CC算法和本文算法），其性能總體上要優(yōu)于基于假設(shè)的算法（灰度世界法、完美反射法）。

表2 不同算法結(jié)果比較

4.5 與普通U-Net網(wǎng)絡(luò)的比較

為了比較本文網(wǎng)絡(luò)與普通U-Net網(wǎng)絡(luò)在顏色校正上的性能，我們用數(shù)據(jù)集Set1對普通U-Net網(wǎng)絡(luò)進(jìn)行了訓(xùn)練和測試，數(shù)據(jù)集的劃分和其它參數(shù)同3.1節(jié)和3.2節(jié)所述。表3所示為本文算法與由普通U-Net網(wǎng)絡(luò)校正顏色后的圖像質(zhì)量，圖5所示為部分由兩種算法校正顏色后的圖像對比。由表3和圖5可以看出，由本文算法恢復(fù)的圖像質(zhì)量更好。

表3 不同算法結(jié)果比較

圖5 不同算法顏色校正效果比較

5 結(jié)語

本文設(shè)計(jì)了一個基于改進(jìn)的U-Net網(wǎng)絡(luò)的端到端的顏色校正神經(jīng)網(wǎng)絡(luò)模型，該模型可以自動預(yù)測輸入圖像的色偏，進(jìn)而校正顏色。通過與現(xiàn)有方法比較，結(jié)果表明，本文方法取得了更好的顏色校正效果，可以顯著改善圖像色偏。下一步的工作是進(jìn)一步優(yōu)化模型和使用更多的圖像進(jìn)行訓(xùn)練和測試，來進(jìn)一步提高恢復(fù)圖像質(zhì)量。