為修復(fù)一張老照片 動用了四塊1080Ti顯卡

嗨！各位，我是一位來自Mail.ru Group計算機(jī)視覺團(tuán)隊的研發(fā)工程師，在這篇文章當(dāng)中，我將講述我們?nèi)绾螢槔宪娙说恼掌瑒?chuàng)造一個基于AI技術(shù)的照片修復(fù)項目。什么是照片修復(fù)呢？它由以下三個步驟組成：

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

-我們基于所發(fā)現(xiàn)的照片缺陷周圍的像素值來進(jìn)行圖像修復(fù)

-我們?yōu)閳D像上色

接下來，我將闡釋照片修復(fù)的每一個步驟，并且告訴你我們?nèi)绾潍@得數(shù)據(jù)，用哪種網(wǎng)絡(luò)訓(xùn)練，取得了哪些成就，以及犯了什么錯誤。

尋找缺陷

我們需要在一張已經(jīng)上傳的照片當(dāng)中找到所有與缺陷相關(guān)的像素值。首先，我們需要弄清楚人們會上傳什么樣的照片。我們與immortal regiment項目的創(chuàng)始人交流過，一個存儲二戰(zhàn)遺留照片的非商業(yè)組織，他們分享數(shù)據(jù)給我們。在此基礎(chǔ)上進(jìn)行分析，我們意識到人們上傳的大部分個人或集體肖像存在中等到大量的缺陷。

接著我們必須收集一個訓(xùn)練集，這個訓(xùn)練集需要對圖像進(jìn)行分割，得到一張圖片，和一張把所有缺陷都標(biāo)注出來的黑底白紋蒙版。做這一步最簡單的方法是讓評估員創(chuàng)立分割蒙版。當(dāng)然，一般人非常清楚怎樣找到缺陷，但那會耗費(fèi)太長時間。

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

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

在我們?yōu)榱隧椖咳蝿?wù)使用這個分割方法時，什么樣的問題會出現(xiàn)呢？

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

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

如果我們通過壓縮處理小圖片，比如說到256或512像素，小瑕疵會因?yàn)閮?nèi)插而消失。然而，我們還需要處理大圖像。在實(shí)際應(yīng)用中我們一般把1024像素的照片瑕疵分割。這就是為什么我們必須將網(wǎng)絡(luò)訓(xùn)練成為適于處理大圖像。但這會引起在單個GPU上處理規(guī)格小批量的問題。

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

最后，基于增加bce的權(quán)重，改變體系結(jié)構(gòu)和使用In-place BatchNorm，我們將分割做得更好，但為了做得更好而使用Test Time Augmentation并不會花費(fèi)許多。我們可以在一張輸入的圖片上運(yùn)行網(wǎng)絡(luò)，獲取鏡像，再重新運(yùn)行網(wǎng)絡(luò)去找到所有的小缺陷。

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

圖像修復(fù)

我們使用了在圖像分割領(lǐng)域很火的Unet網(wǎng)絡(luò)。為了做圖像修復(fù)，我們需要上傳原始圖像和一個掩模(使用1標(biāo)記了所有干凈區(qū)域的像素，使用0標(biāo)記了需要好修復(fù)的像素)。數(shù)據(jù)是這樣被搜集的：在一些開源的圖像數(shù)據(jù)集上(例如：OpenImagesV4)加入了一些在生活中會出現(xiàn)的照片瑕疵，然后我們訓(xùn)練網(wǎng)絡(luò)去重新修復(fù)有瑕疵的這部分。

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

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

我們花了5天來訓(xùn)練這個網(wǎng)絡(luò)。在最后一天我們停止了BatchNorms 來使得受損部分的邊緣更不易被看見。

這個網(wǎng)絡(luò)處理一副512*512的圖像需要花費(fèi)50ms，驗(yàn)證集上的PSNR（“Peak Signal to Noise Ratio”的縮寫，即峰值信噪比）等于26.4。然而在這種任務(wù)中，你不能完全依賴這個指標(biāo)。為了選擇最好的模型，我們在一些驗(yàn)證圖片上運(yùn)行幾個好的模型，將結(jié)果隱藏，然后投票選出我們最喜歡的模型修復(fù)過的圖片，這是我們?nèi)绾芜x擇最終模型的方法。

我之前提到過我們?nèi)斯ぴ诟蓛舻膱D片上加了一些缺陷 。所以在訓(xùn)練時需要一直跟蹤添加的缺陷的最大面積。以防當(dāng)你"喂"給網(wǎng)絡(luò)一張它從沒有在訓(xùn)練中處理過的，有很大缺陷的圖片。網(wǎng)絡(luò)可能會失控并且產(chǎn)生不可用的結(jié)果。因此，如果你想修復(fù)有大缺陷的圖片，增大你訓(xùn)練集里的缺陷。

這是一個關(guān)于我們的算法效果的例子：

著色

我們已經(jīng)將圖片有缺陷的部分分割了出來并且修復(fù)了他們；第三步—;—;顏色復(fù)原。就像我之前說的，Immortal Regiment項目有很多個人或集體的相片。我們希望網(wǎng)絡(luò)可以在修復(fù)這些照片上起到很好的作用。因?yàn)楝F(xiàn)有的服務(wù)不能有效且快的著色，所以我們決定提出我們自己的著色方案。我們希望我們的被著色的圖片更可信。

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

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

為什么要選擇YCbCr色彩表達(dá)式？因?yàn)槿搜蹖γ靼底兓拿舾谐潭葧葘ι首兓拿舾谐潭雀?。這是我們選擇復(fù)用了人眼更敏感的Y通道（亮度）及預(yù)測我們可能出錯的Cb和Cr通道值的原因，因?yàn)槲覀儾粫茏⒁忸伾e誤。這個具體的特性廣泛用于彩色電視產(chǎn)生的初期，那時通道傳輸能力還不足以傳輸所有顏色。圖片是依據(jù)YCbCr色彩表達(dá)式來傳輸?shù)模琘通道數(shù)值沒有改變，Cb和Cr的數(shù)值均減半。

如何新建一個基準(zhǔn)模型

我們將現(xiàn)有的CbCr值輸入進(jìn)帶有預(yù)訓(xùn)練好編碼器的Unet網(wǎng)絡(luò)，然后最小化L1損失函數(shù)來預(yù)測合理的CbCr值。我們希望為照片上色，因此除了OpenImages數(shù)據(jù)集的相片之外，我們需要更多任務(wù)特定的相片。

我們?nèi)ツ睦锏玫酱┲娧b的人的彩色照呢？在網(wǎng)上會有些人因?yàn)榕d趣或者金錢為老照片上色。他們做的很仔細(xì)很準(zhǔn)確。他們是根據(jù)檔案材料來為制服，肩牌和勛章上色，所以他們的成果很值得信賴。最后，我們使用了200幅手工上色的身著軍裝的人的照片。

另一個有用的數(shù)據(jù)源是The Workers’ and Peasants’ Red Army 網(wǎng)站。它的創(chuàng)辦者之一幾乎拍了所有自己身著二戰(zhàn)時期軍裝的照片。

在某些照片中，他模仿了著名照片人物的姿勢。他的照片都是白色背景，這可以讓我們在背景上添加不同的自然風(fēng)景來擴(kuò)充數(shù)據(jù)。我們也使用了一些平常的相片，并在上面添加了肩章和其他戰(zhàn)爭相關(guān)的元素。

我們訓(xùn)練了 AlbuNet-50—;—;這是一個使用了預(yù)訓(xùn)練好的ResNet-50 作為編碼器的Unet。在預(yù)測時，這個網(wǎng)絡(luò)給出了一些合理的結(jié)果：皮膚是粉色的，眼睛—;—;灰綠色，肩章—;—;微黃色。然而，會有一些區(qū)域沒有被著色。這是因?yàn)楦ヮA(yù)測一些別的顏色相比，不做任何變化會使L1誤差有更優(yōu)的結(jié)果。

我們正在比較我們的結(jié)果與原始真實(shí)照片-由Klimbim手工著色

我們該如何解決這個問題呢？我們需要一個判別：一個輸入圖像，并判斷圖像是否是真實(shí)的圖像的神經(jīng)網(wǎng)絡(luò)。下面的其中一張圖片是手工著色的，另一張是由我們的生成器AlbuNet-50繪制的。人類如何區(qū)分手動和自動著色的照片呢?通過觀察細(xì)節(jié)。你能告訴我們的哪個照片是基準(zhǔn)解決方案得到的嗎？

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

我們使用了基于自注意機(jī)制的GAN的論文中的判別器。這是一個很小的卷積網(wǎng)絡(luò)，在最上面的幾層是被稱為自注意力的層。它使我們的網(wǎng)絡(luò)“更加關(guān)注”圖像細(xì)節(jié)。我們也使用譜歸一化。你可以在上面的論文中找到更多的細(xì)節(jié)。我們利用L1損失項和判別器損失項的組合訓(xùn)練了網(wǎng)絡(luò)?，F(xiàn)在得到的網(wǎng)絡(luò)在圖像的細(xì)節(jié)上著色更好，背景看起來更一致。再舉一個例子:左邊是只訓(xùn)練L1損失項的網(wǎng)絡(luò)得到的結(jié)果;右邊是結(jié)合L1損失項的結(jié)果。

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

當(dāng)我們已經(jīng)創(chuàng)建了一個生產(chǎn)系統(tǒng)，并推出了一個網(wǎng)站，我們繼續(xù)驗(yàn)證，獲得結(jié)果，我們最好盡量減少每個像素L1損失，而不是減少感知損失。為了計算它，我們將網(wǎng)絡(luò)的預(yù)測和一張ground-truthl照片提供給VGG-16網(wǎng)絡(luò)，獲取底層的特征map，并利用最小均方誤差進(jìn)行比較。這種方法修復(fù)了更多的區(qū)域，并給出了更有色調(diào)的結(jié)果。

簡要回顧

Unet是一個非常棒的模型。在第一個分割任務(wù)中，我們在訓(xùn)練過程中遇到了一個問題，就是處理高分辨率的圖像，這就是為什么我們使用In-Place 批歸一化。在我們的第二個任務(wù)(圖像修復(fù))中，我們使用了部分卷積而不是標(biāo)準(zhǔn)卷積，這讓我們得到了更好的結(jié)果。在進(jìn)行著色時，我們增加了一個簡單的判別器網(wǎng)絡(luò)，它可以對生成不真實(shí)圖像的生成器進(jìn)行懲罰。我們還使用了感知損失。

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

我們從War Album project 中獲得了一些照片，并在這些神經(jīng)網(wǎng)絡(luò)上進(jìn)行了處理。以下是我們得到的結(jié)果：

此外，你還可以更詳細(xì)地查看原始圖像和所有的處理過程，詳情訪問：點(diǎn)此。