修復一張老照片 動用四塊1080Ti顯卡

嗨！各位，我是一位來自Mail.ru Group計算機視覺團隊的研發工程師，在這篇文章當中，我將講述我們如何為老軍人的照片創造一個基於AI技術的照片修復項目。什麼是照片修復呢？它由以下三個步驟組成：

-我們找到所有的照片缺陷：摺痕，磨損，破洞

-我們基於所發現的照片缺陷周圍的像素值來進行圖像修復

-我們為圖像上色

接下來，我將闡釋照片修復的每一個步驟，並且告訴你我們如何獲得數據，用哪種網絡訓練，取得了哪些成就，以及犯了什麼錯誤。

尋找缺陷

我們需要在一張已經上傳的照片當中找到所有與缺陷相關的像素值。首先，我們需要弄清楚人們會上傳什麼樣的照片。我們與immortal regiment項目的創始人交流過，一個存儲二戰遺留照片的非商業組織，他們分享數據給我們。在此基礎上進行分析，我們意識到人們上傳的大部分個人或集體肖像存在中等到大量的缺陷。

接着我們必須收集一個訓練集，這個訓練集需要對圖像進行分割，得到一張圖片，和一張把所有缺陷都標註出來的黑底白紋蒙版。做這一步最簡單的方法是讓評估員創立分割蒙版。當然，一般人非常清楚怎樣找到缺陷，但那會耗費太長時間。

標記一張照片中缺陷處的像素將會花費一個小時或一整個工作日，因此，在幾周內收集一個超過100張圖片的訓練集是不簡單的。這就是為什麼我們嘗試加強我們的數據，並創造我們自己的缺陷：拍一張照片，用圖片上的隨機紋理增加缺陷，最終得到一張顯示有缺陷的圖像部分的蒙版。在沒有增強的情況下，我們已經獲得了含有68張手動標記照片的訓練集和和11張照片的驗證集。

最流行的分割方法是：使用unet和預訓練編碼器，將BCE（binary cross-entropy）與DICE（S?rensen–Dice coefficient）的和最小化。

在我們為了項目任務使用這個分割方法時，什麼樣的問題會出現呢？

雖然照片看起來像有許多的缺陷，非常老舊而殘破，有缺陷的部分仍然遠遠小於沒有受到損傷的部分。為了解決這個問題，我們可以增加BCE中的積極類內權重，最優權重應是清晰像素與缺陷像素的比率。

第二個問題是如果我們用可以立刻上手使用的預訓練編碼器和UNET，將會損失許多位置數據。第1層如果由一個核為5、步長為2的卷積組成，網絡運行得更快。我們犧牲網絡運行的時間來換取更好的缺陷定位。我們在第1層之後移除最大的pooling層，將卷積步長減至1，卷積核減至3。

如果我們通過壓縮處理小圖片，比如説到256或512像素，小瑕疵會因為內插而消失。然而，我們還需要處理大圖像。在實際應用中我們一般把1024像素的照片瑕疵分割。這就是為什麼我們必須將網絡訓練成為適於處理大圖像。但這會引起在單個GPU上處理規格小批量的問題。

在訓練中，我們我們可以在一個GPU上匹配20張左右的圖片。因此我們最終會在BatchNorm層得到不準確的平均值和標準偏差值。我們可以通過 In-place BatchNorm 來解決這個問題。一方面節省內存空間，另一方面擁有一個 Synchronized BatchNorm版本，可以同步所有GPU上的數據。這樣我們計算平均值和標準偏差值就不是參照一個GPU上的20張圖片，而是4個GPU上的80張圖片，這改善了網絡集成。

最後，基於增加bce的權重，改變體系結構和使用In-place BatchNorm，我們將分割做得更好，但為了做得更好而使用Test Time Augmentation並不會花費許多。我們可以在一張輸入的圖片上運行網絡，獲取鏡像，再重新運行網絡去找到所有的小缺陷。

用了四個GeForce 1080Ti顯卡，花費了18個小時，相對時間較長，但這就是精益求精的代價。

圖像修復

我們使用了在圖像分割領域很火的Unet網絡。為了做圖像修復，我們需要上傳原始圖像和一個掩模(使用1標記了所有乾淨區域的像素，使用0標記了需要好修復的像素)。數據是這樣被蒐集的：在一些開源的圖像數據集上(例如：OpenImagesV4)加入了一些在生活中會出現的照片瑕疵，然後我們訓練網絡去重新修復有瑕疵的這部分。

我們是如何使用在圖像修復中使用Unet的呢？

我們使用了部分卷積替代了原始的那一版。當我們使用一些核去卷積一個區域時，我們沒有把損壞的像素算在內。這使得圖像修復更精密。從英偉達的論文來舉個例子，他們使用了Unet結合默認的2維卷積的修復結果在中間這幅圖，使用部分卷積的修復結果在右邊這幅圖。

我們花了5天來訓練這個網絡。在最後一天我們停止了BatchNorms 來使得受損部分的邊緣更不易被看見。

這個網絡處理一副512*512的圖像需要花費50ms，驗證集上的PSNR（“Peak Signal to Noise Ratio”的縮寫，即峯值信噪比）等於26.4。然而在這種任務中，你不能完全依賴這個指標。為了選擇最好的模型，我們在一些驗證圖片上運行幾個好的模型，將結果隱藏，然後投票選出我們最喜歡的模型修復過的圖片，這是我們如何選擇最終模型的方法。

我之前提到過我們人工在乾淨的圖片上加了一些缺陷  。所以在訓練時需要一直跟蹤添加的缺陷的最大面積。以防當你"喂"給網絡一張它從沒有在訓練中處理過的，有很大缺陷的圖片。網絡可能會失控並且產生不可用的結果。因此，如果你想修復有大缺陷的圖片，增大你訓練集裏的缺陷。

這是一個關於我們的算法效果的例子：

着色

我們已經將圖片有缺陷的部分分割了出來並且修復了他們；第三步——顏色復原。就像我之前説的，Immortal Regiment項目有很多個人或集體的相片。我們希望網絡可以在修復這些照片上起到很好的作用。因為現有的服務不能有效且快的着色，所以我們決定提出我們自己的着色方案。我們希望我們的被着色的圖片更可信。

github有個很流行的相片着色repo。它有着很好的效果但是還是有很多問題。例如，它有可能會將衣服塗成藍色，這也是我們沒有使用他的原因。

所以，我們決定在圖像着色方面提出一個算法。其中最簡單的方案是：預測一張黑白照片的R,G,B三個通道的值。但是，我們可以使工作更簡單：使用YCbCr色彩表達式來替換RGB色彩表達式。Y通道表示亮度。一張上傳的黑白照片只在Y通道里有數值，我們會複用這個數值。而只需要預測Cb和Cr通道的值。Cb是藍色與亮度的差值，Cr是紅色與亮度的差值。

為什麼要選擇YCbCr色彩表達式？因為人眼對明暗變化的敏感程度會比對色彩變化的敏感程度更高。這是我們選擇複用了人眼更敏感的Y通道（亮度）及預測我們可能出錯的Cb和Cr通道值的原因，因為我們不會很注意顏色錯誤。這個具體的特性廣泛用於彩色電視產生的初期，那時通道傳輸能力還不足以傳輸所有顏色。圖片是依據YCbCr色彩表達式來傳輸的，Y通道數值沒有改變，Cb和Cr的數值均減半。

如何新建一個基準模型

我們將現有的CbCr值輸入進帶有預訓練好編碼器的Unet網絡，然後最小化L1損失函數來預測合理的CbCr值。我們希望為照片上色，因此除了OpenImages數據集的相片之外，我們需要更多任務特定的相片。

我們去哪裏得到穿着軍裝的人的彩色照呢？在網上會有些人因為興趣或者金錢為老照片上色。他們做的很仔細很準確。他們是根據檔案材料來為制服，肩牌和勳章上色，所以他們的成果很值得信賴。最後，我們使用了200幅手工上色的身着軍裝的人的照片。

另一個有用的數據源是The Workers’ and Peasants’ Red Army 網站。它的創辦者之一幾乎拍了所有自己身着二戰時期軍裝的照片。

在某些照片中，他模仿了著名照片人物的姿勢。他的照片都是白色背景，這可以讓我們在背景上添加不同的自然風景來擴充數據。我們也使用了一些平常的相片，並在上面添加了肩章和其他戰爭相關的元素。

我們訓練了 AlbuNet-50——這是一個使用了預訓練好的ResNet-50 作為編碼器的Unet。在預測時，這個網絡給出了一些合理的結果：皮膚是粉色的，眼睛——灰綠色，肩章——微黃色。然而，會有一些區域沒有被着色。這是因為跟去預測一些別的顏色相比，不做任何變化會使L1誤差有更優的結果。

我們正在比較我們的結果與原始真實照片-由Klimbim手工着色

我們該如何解決這個問題呢？我們需要一個判別：一個輸入圖像，並判斷圖像是否是真實的圖像的神經網絡。下面的其中一張圖片是手工着色的，另一張是由我們的生成器AlbuNet-50繪製的。人類如何區分手動和自動着色的照片呢?通過觀察細節。你能告訴我們的哪個照片是基準解決方案得到的嗎？

回答：左邊的圖片是手動上色的，右邊是自動上色的。

我們使用了基於自注意機制的GAN的論文中的判別器。這是一個很小的卷積網絡，在最上面的幾層是被稱為自注意力的層。它使我們的網絡“更加關注”圖像細節。我們也使用譜歸一化。你可以在上面的論文中找到更多的細節。我們利用L1損失項和判別器損失項的組合訓練了網絡。現在得到的網絡在圖像的細節上着色更好，背景看起來更一致。再舉一個例子:左邊是隻訓練L1損失項的網絡得到的結果;右邊是結合L1損失項的結果。

在四個GeForce 1080Ti上，訓練過程花費了兩天時間。處理一張512x512的圖片需要30毫秒。驗證集的最小均方誤差(MSE)為34.4。就像圖像修復一樣，基於我們並不想選擇的評價準則。這就是為什麼我們選擇了6個具有最佳驗證指標的模型，並直接為最佳模型投票。

當我們已經創建了一個生產系統，並推出了一個網站，我們繼續驗證，獲得結果，我們最好儘量減少每個像素L1損失，而不是減少感知損失。為了計算它，我們將網絡的預測和一張ground-truthl照片提供給VGG-16網絡，獲取底層的特徵map，並利用最小均方誤差進行比較。這種方法修復了更多的區域，並給出了更有色調的結果。

簡要回顧

Unet是一個非常棒的模型。在第一個分割任務中，我們在訓練過程中遇到了一個問題，就是處理高分辨率的圖像，這就是為什麼我們使用In-Place 批歸一化。在我們的第二個任務(圖像修復)中，我們使用了部分卷積而不是標準卷積，這讓我們得到了更好的結果。在進行着色時，我們增加了一個簡單的判別器網絡，它可以對生成不真實圖像的生成器進行懲罰。我們還使用了感知損失。

第二個結論——評估人員是必不可少的。不僅在創建分割掩碼階段，而且在最終的結果驗證階段。最後，我們給用户提供了三張照片:一幅帶有缺陷的原始照片，一幅帶有缺陷的彩色照片，以及一幅簡單的彩色照片，以防缺陷搜索和圖像修復算法出錯。

我們從War Album project 中獲得了一些照片，並在這些神經網絡上進行了處理。以下是我們得到的結果：