基于改進U-Net網絡的圖像自動白平衡算法

全雪峰

（南陽醫學高等專科學校 衛生管理系，河南 南陽 473061）

0 引言

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

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

近年來基于卷積神經網絡的深度學習算法得到快速發展[8-11]。如Shi等人[12]提出用于圖像光照估計DS-Net網絡；Hu等人[13]提出基于置信度加權池的全卷積顏色恒常性算法；Zhu等人[14]對用于圖像風格遷移的網絡結構進行改進，并用于圖像顏色校正；Shen等人[15]根據Retinex顏色恒常性理論提出MSR-net網絡，用于解決低光照下顏色增強問題；Simone等人[16]提出一種準無監督的深度卷積神經網絡顏色恒常性算法。本文設計了一種新的端到端的圖像顏色校正網絡。實驗結果表明，本文算法可以有效校正圖像色偏。

1 U-Net網絡介紹

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

2 基于改進U-Net網絡的顏色校正方法

2.1 網絡結構

本文對U-Net網絡進行改進，設計了一種新的用于顏色校正的端到端的網絡結構，如圖1所示。

圖1 網絡結構

由圖1可見，整個網絡主要有由三部分組成：（1）用于提取特征的編碼器；（2）用于提取圖像細節信息的稠密連接塊（Dense Block）；（3）用于恢復圖像的解碼器。整個網絡類似于U-Net網絡結構，

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

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

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

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

2.2 損失函數

本文在訓練網絡時使用的損失函數是均方誤差（Mean Squared Error，MSE）[14]，定義如下：

式中，Io是網絡輸出圖像，Ir是參考圖像，P是圖像塊，N是像素總數。

3 實驗

本文的實驗平臺為Deepin Linux 11.5，CPU為intel i5 7600，GPU為Nivida GForce GTX 1050Ti，16G內存，采用Pytorch深度學習框架實現。

3.1 數據集

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

3.2 訓練參數

考慮到將學習速率設置過大，難以找到最優解，設置過小，網絡收斂緩慢，本文設置初始學習率為0.01，權值衰減為0.0001，batch=8，采用Adam優化器。由實驗發現，經過10個epoch后就可以充分訓練網絡，如圖2所示。

圖2 Epoch-loss曲線

4 結果分析

4.1 評價指標

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

4.2 不同卷積設計對顏色校正效果的影響

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

由圖3和表1可知，當網絡設計為{3, 2}時，顏色校正后的圖像質量最好，原因是網絡采集到的圖像特征最多。在隨后的實驗中，采用的都是編碼器和解碼器的所有層的卷積核大小為3，每層有2個卷積。

圖3 不同參數圖像校正結果

表1 不同參數圖像校正結果

4.3 不同算法的視覺比較

圖4所示為部分色偏圖像用灰度世界法[1]、Quasi-U CC算法[16]和本文算法校正顏色后的視覺比較。由圖4可以看出，除了圖片f以外，灰度世界法的顏色校正效果均不理想，尤其是對圖片d和圖片e的校正結果出現明顯的色偏，這是因為該算法假設在標準光源下的場景，其平均反射是無色差的，這樣當場景包含有大片單色區域時，校正的結果會出現色偏。Quasi-U CC算法的顏色校正效果優于灰度世界法，但對圖片d，基本未校正顏色，對于圖片g，校正后的圖片較暗。本文算法的顏色校正結果均取得了較好的視覺效果，優于灰度世界法和Quasi-U CC算法。

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

4.4 不同算法客觀比較

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

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

表2 不同算法結果比較

4.5 與普通U-Net網絡的比較

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

表3 不同算法結果比較

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

5 結語

本文設計了一個基于改進的U-Net網絡的端到端的顏色校正神經網絡模型，該模型可以自動預測輸入圖像的色偏，進而校正顏色。通過與現有方法比較，結果表明，本文方法取得了更好的顏色校正效果，可以顯著改善圖像色偏。下一步的工作是進一步優化模型和使用更多的圖像進行訓練和測試，來進一步提高恢復圖像質量。