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

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

-我們基于所發(fā)現(xiàn)的照片缺陷周?chē)南袼刂祦?lái)進(jìn)行圖像修復(fù)

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

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

尋找缺陷

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

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

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

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

在我們?yōu)榱隧?xiàng)目任務(wù)使用這個(gè)分割方法時(shí)，什么樣的問(wèn)題會(huì)出現(xiàn)呢？

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

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

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

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

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

用了四個(gè)GeForce 1080Ti顯卡，花費(fèi)了18個(gè)小時(shí)，相對(duì)時(shí)間較長(zhǎng)，但這就是精益求精的代價(jià)。

圖像修復(fù)

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

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

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

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

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

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

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

著色

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

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

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

為什么要選擇YCbCr色彩表達(dá)式？因?yàn)槿搜蹖?duì)明暗變化的敏感程度會(huì)比對(duì)色彩變化的敏感程度更高。這是我們選擇復(fù)用了人眼更敏感的Y通道（亮度）及預(yù)測(cè)我們可能出錯(cuò)的Cb和Cr通道值的原因，因?yàn)槲覀儾粫?huì)很注意顏色錯(cuò)誤。這個(gè)具體的特性廣泛用于彩色電視產(chǎn)生的初期，那時(shí)通道傳輸能力還不足以傳輸所有顏色。圖片是依據(jù)YCbCr色彩表達(dá)式來(lái)傳輸?shù)?，Y通道數(shù)值沒(méi)有改變，Cb和Cr的數(shù)值均減半。

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

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

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

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

在某些照片中，他模仿了著名照片人物的姿勢(shì)。他的照片都是白色背景，這可以讓我們?cè)诒尘吧咸砑硬煌淖匀伙L(fēng)景來(lái)擴(kuò)充數(shù)據(jù)。我們也使用了一些平常的相片，并在上面添加了肩章和其他戰(zhàn)爭(zhēng)相關(guān)的元素。

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

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

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

回答：左邊的圖片是手動(dòng)上色的，右邊是自動(dòng)上色的。

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

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

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

簡(jiǎn)要回顧

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

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

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

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