AI-Based Photo Restoration

大家好！我是mail.ru集團(tuán)計(jì)算機(jī)視覺團(tuán)隊(duì)的研究工程師。在這篇文章中，我將講述一個(gè)關(guān)于我們?nèi)绾螢榕f的軍事照片創(chuàng)建基于人工智能的照片修復(fù)項(xiàng)目。什么是“照片復(fù)原”？它包括三個(gè)步驟：
?發(fā)現(xiàn)所有圖像的缺陷：破裂、刮傷、孔洞；
?根據(jù)周圍的像素值修復(fù)發(fā)現(xiàn)的缺陷；
?給圖像上色。
此外，我將描述照片恢復(fù)的每一步，并告訴你我們?nèi)绾潍@得數(shù)據(jù)，我們訓(xùn)練了哪些網(wǎng)絡(luò)，我們完成了什么，以及我們犯了什么錯(cuò)誤。

尋找缺陷

我們想在上傳的照片中找到所有與缺陷相關(guān)的像素。首先，我們需要弄清楚人們會(huì)上傳什么樣的圖片。我們采訪了“不朽軍團(tuán)”項(xiàng)目的創(chuàng)始人，這是一個(gè)非商業(yè)組織，存儲(chǔ)了二戰(zhàn)時(shí)遺留的照片，他們與我們分享了他們的數(shù)據(jù)。經(jīng)過分析，我們注意到人們上傳的大多是個(gè)人或團(tuán)體肖像，存在中等到大量的缺陷。
然后我們得收集一套訓(xùn)練器材。分割任務(wù)的訓(xùn)練集是一個(gè)圖像和一個(gè)標(biāo)記所有缺陷的掩碼。最簡(jiǎn)單的方法是讓評(píng)估人員創(chuàng)建分段掩碼。當(dāng)然，人們很清楚如何發(fā)現(xiàn)缺陷，但這需要很長(zhǎng)時(shí)間。

在一張照片上標(biāo)記缺陷像素可能需要一個(gè)小時(shí)或整個(gè)工作日。因此，要在幾周內(nèi)收集到超過100張圖片的訓(xùn)練集并不容易。這就是為什么我們?cè)噲D增加我們的數(shù)據(jù)并創(chuàng)建我們自己的缺陷：我們會(huì)拍一張好的照片，在圖像上使用隨機(jī)游動(dòng)添加缺陷，最后得到一個(gè)顯示有缺陷的圖像部分的遮罩。在沒有增強(qiáng)的情況下，我們?cè)谟?xùn)練集中有68張手動(dòng)標(biāo)記的照片，在驗(yàn)證集中有11張。

最流行的分割方法是：使用帶預(yù)訓(xùn)練編碼器的Unet，并最小化BCE（二進(jìn)制交叉熵）和DICE（binary cross-entropy]) 。

當(dāng)我們?cè)谌蝿?wù)中使用這種分割方法時(shí)會(huì)出現(xiàn)什么問題？

?即使照片看起來有很多缺陷，很舊很破舊，但有缺陷的區(qū)域仍然比未損壞的小得多。為了解決這個(gè)問題，我們可以增加BCE中的正類權(quán)重，一個(gè)最佳的權(quán)重是干凈像素與有缺陷像素的比率。
?第二個(gè)問題是，如果我們使用帶預(yù)訓(xùn)練編碼器（例如Albunet-18）的開箱即用Unet，我們會(huì)丟失很多位置數(shù)據(jù)。Albunet-18的第一層由一個(gè)核5的卷積和一個(gè)等于2的步長(zhǎng)組成。它可以讓網(wǎng)絡(luò)快速工作。為了更好地定位缺陷，我們交換了網(wǎng)絡(luò)操作時(shí)間：刪除了第一層之后的最大池，將步長(zhǎng)減少到1，將卷積核減少到3。
?如果我們通過壓縮小圖像（例如，將其壓縮到256 x 256或512 x 512像素）來處理小圖像，則由于插值，小缺陷將消失。因此，我們需要處理更大的圖像。我們目前正在生產(chǎn)中對(duì)1024x 1024大小的照片進(jìn)行缺陷分割。這就是為什么我們必須在大圖像上訓(xùn)練網(wǎng)絡(luò)。但是，這會(huì)導(dǎo)致單個(gè)GPU上的小批量出現(xiàn)問題。
?在培訓(xùn)期間，我們可以在一個(gè)GPU上安裝大約20個(gè)圖像。因此，我們?cè)赽atchnorm層中的平均值和標(biāo)準(zhǔn)偏差值不準(zhǔn)確。我們可以使用就地batchnorm解決這個(gè)問題，一方面節(jié)省內(nèi)存空間，另一方面，它有一個(gè)同步的batchnorm版本，可以同步所有GPU上的統(tǒng)計(jì)信息。現(xiàn)在，我們計(jì)算的平均值和標(biāo)準(zhǔn)偏差值不是針對(duì)單個(gè)GPU上的20個(gè)圖像，而是針對(duì)4個(gè)GPU上的80個(gè)圖像。這改善了網(wǎng)絡(luò)收斂性。
最后，通過增加bce權(quán)重、改變體系結(jié)構(gòu)、使用原位batchnorm，使分割效果更好。然而，通過增加測(cè)試時(shí)間來做更好的事情不會(huì)花費(fèi)太多的成本。我們可以在輸入圖片上運(yùn)行一次網(wǎng)絡(luò)，然后對(duì)其進(jìn)行鏡像并重新運(yùn)行網(wǎng)絡(luò)以查找所有的小缺陷。

網(wǎng)絡(luò)在18小時(shí)內(nèi)匯聚在4個(gè)Geforce 1080ti上。推理需要290毫秒。這很長(zhǎng)，但這是我們優(yōu)于默認(rèn)性能的代價(jià)。Validation DICE equals 0,35, and ROCAUC — 0,93.

圖像修復(fù)

與分割任務(wù)一樣，我們使用Unet。要進(jìn)行修補(bǔ)，我們需要上傳一張?jiān)紙D像和一個(gè)遮罩，在那里我們用1和0標(biāo)記所有清潔區(qū)域-所有我們想要修補(bǔ)的像素。這就是我們收集數(shù)據(jù)的方式：對(duì)于來自開源圖像數(shù)據(jù)集（例如openimagesv4）的任何照片，我們都添加了與現(xiàn)實(shí)生活中看到的類似的缺陷。然后我們訓(xùn)練網(wǎng)絡(luò)恢復(fù)丟失的部分。
我們?nèi)绾涡薷倪@個(gè)任務(wù)的Unet？
我們可以用部分卷積代替原來的卷積。其思想是，當(dāng)我們用一些核卷積一個(gè)區(qū)域時(shí)，我們不考慮缺陷像素值,這使得修補(bǔ)更加精確。我們向您展示了最近Nvidia論文中的一個(gè)例子。他們?cè)谥虚g的圖片中使用了默認(rèn)的二維卷積的Unet，在右邊的圖片中使用了部分卷積的Unet。

我們訓(xùn)練花了五天。最后一天，我們凍結(jié)了批處理規(guī)范，以減少著色部分的邊界。
網(wǎng)絡(luò)處理一張512 x 512的圖片需要50毫秒。驗(yàn)證PSNR等于26,4。但是，您不能完全依賴此任務(wù)中的度量。為了選擇最好的模型，我們對(duì)估值圖像運(yùn)行了幾個(gè)好的模型，匿名投票給我們最喜歡的模型，這就是我們選擇最終模型的方式。
如我之前提到過，我們?nèi)藶榈卦诟蓛舻膱D像中添加了一些缺陷。在訓(xùn)練過程中，您應(yīng)該始終跟蹤添加的缺陷的最大大小；如果您將一個(gè)具有非常大缺陷的圖像饋送到網(wǎng)絡(luò)，而在訓(xùn)練階段它從未處理過該圖像，則網(wǎng)絡(luò)將瘋狂運(yùn)行并產(chǎn)生不適用的結(jié)果。因此，如果你需要修復(fù)大的缺陷，用它們來增加你的訓(xùn)練集。
下面是我們的算法如何工作的示例：

著色

我們將缺陷分割并修補(bǔ)；第三步-顏色重建。就像我之前說的，照片中有很多個(gè)人和團(tuán)體的肖像。我們希望我們的網(wǎng)絡(luò)能和他們很好地配合。我們決定提出自己的彩色化，因?yàn)闆]有一個(gè)現(xiàn)有的服務(wù)可以快速和有效的彩色肖像。我們希望彩色照片更真實(shí)可信。

Github有一個(gè)流行的照片著色庫(kù)。它做得很好也有一些問題。例如，它傾向于把衣服涂成藍(lán)色。所以我不用它。
因此，我們決定創(chuàng)建一個(gè)圖像著色算法。最明顯的想法是：拍攝一張黑白圖像，預(yù)測(cè)三個(gè)通道：紅色、綠色和藍(lán)色。然而，我們可以使我們的工作更容易：工作不是用RGB顏色表示，而是與YCbCr顏色表示。y分量是亮度（luma）。上傳的黑白圖像是y通道，我們將重用它。現(xiàn)在我們需要預(yù)測(cè)Cb和Cr：Cb是藍(lán)色和亮度的差，Cr是紅色和亮度的差。

為什么我們選擇YCbCr代表？人眼對(duì)亮度的變化比對(duì)顏色的變化更敏感。這就是為什么我們重復(fù)使用人眼最敏感的Y分量（亮度），并預(yù)測(cè)Cb和Cr，因?yàn)槲覀儾荒芎芎玫刈⒁獾筋伾奶摷傩裕晕覀兛赡軙?huì)犯錯(cuò)誤。這一特性在彩色電視出現(xiàn)之初就得到了廣泛的應(yīng)用，當(dāng)時(shí)的頻道容量還不足以傳輸所有的顏色。圖像在YCbCr中傳輸，不改變y分量，Cb和Cr減少一半。

如何創(chuàng)建基線

我們可以采用帶預(yù)訓(xùn)練編碼器的Unet，并將現(xiàn)有CbCr值和預(yù)測(cè)值之間的l1損失最小化。我們想給人像上色，因此，除了openimages照片，我們還需要更多特定于任務(wù)的照片。
我們?cè)谀睦锟梢缘玫酱┲娧b的人的彩色照片？互聯(lián)網(wǎng)上有人把舊照片上色當(dāng)作一種愛好或是一種代價(jià)。他們做得很仔細(xì)，盡量精確。當(dāng)他們給制服、肩章和獎(jiǎng)?wù)律仙珪r(shí)，他們會(huì)參考檔案材料，因此他們的工作成果是值得信賴的。總而言之，我們使用了200張手工著色的照片，上面有穿著軍裝的人。
另一個(gè)有用的數(shù)據(jù)來源是工農(nóng)紅軍網(wǎng)站。它的一個(gè)創(chuàng)立者幾乎每到第二次世界大戰(zhàn)時(shí)都會(huì)穿著蘇維埃的制服拍照。

在一些照片中，他模仿著名的檔案照片中人們的姿勢(shì)。他的照片有白色背景是一件好事：它讓我們可以通過在背景中添加各種自然物體來很好地?cái)U(kuò)充數(shù)據(jù)。我們還使用了一些常規(guī)的肖像，以徽章和其他戰(zhàn)時(shí)特征作為補(bǔ)充。
我們訓(xùn)練了Albunet-50——它是一個(gè)使用預(yù)先訓(xùn)練的Resnet-50作為編碼器的Unet。網(wǎng)絡(luò)開始提供足夠的結(jié)果：皮膚是粉紅色的，眼睛是灰藍(lán)色的，肩膀是黃色的。然而，問題是它讓照片上的一些區(qū)域沒有受到影響。這是因?yàn)楦鶕?jù)錯(cuò)誤L1，在不做任何事情比試圖預(yù)測(cè)某種顏色更好的地方找到這樣的最佳值。

我們正在將我們的結(jié)果與地面真實(shí)照片進(jìn)行比較-這是由Klimbim完成的手動(dòng)著色。
我們?cè)鯓硬拍芙鉀Q這個(gè)問題？我們需要一個(gè)鑒別器：一個(gè)接收?qǐng)D像并告訴我們圖像是否真實(shí)的神經(jīng)網(wǎng)絡(luò)。下面的一張圖片是手動(dòng)著色的，另一張是由我們的生成器Albunet-50著色的。人類如何區(qū)分手動(dòng)和自動(dòng)彩色照片？通過觀察細(xì)節(jié)。你能告訴我們的基線解決方案后自動(dòng)著色的照片在哪張圖嗎？

答：左邊的圖片是手動(dòng)著色，右邊的是自動(dòng)著色。
我們使用了來自自聚焦Self-Attention GAN鑒別器。這是一個(gè)小型的卷積網(wǎng)絡(luò)，所謂的自我關(guān)注建立在頂層。它允許我們“更加關(guān)注”圖像細(xì)節(jié)，我們也使用光譜標(biāo)準(zhǔn)化。你可以在別的文章中找到更多的信息。我們用L1損失和鑒別器損失的組合來訓(xùn)練網(wǎng)絡(luò)。現(xiàn)在，網(wǎng)絡(luò)使圖像的細(xì)節(jié)顏色更好，背景看起來更加一致。再舉一個(gè)例子：左邊是只訓(xùn)練了L1損耗的網(wǎng)絡(luò)工作；右邊是L1鑒別器損耗的組合。

在四臺(tái)GeForce 1080ti上訓(xùn)練過程花了兩天時(shí)間。處理一張512 x 512的圖片需要30毫秒的時(shí)間。Validation MSE — 34.4。就像修復(fù)一樣，度量你不想依賴于度量。這就是為什么我們選擇了六個(gè)具有最佳驗(yàn)證指標(biāo)的模型，并盲投選出了最佳模型。
當(dāng)我們已經(jīng)創(chuàng)建了一個(gè)制作系統(tǒng)并發(fā)布了一個(gè)網(wǎng)站時(shí)，我們繼續(xù)進(jìn)行實(shí)驗(yàn)并得出結(jié)論，我們最好不是將每像素L1的損失降到最低，而是將感知損失降到最低。為了計(jì)算它，我們將網(wǎng)絡(luò)預(yù)測(cè)和地面實(shí)況照片輸入VGG-16網(wǎng)絡(luò)，獲取底層特征圖并與MSE進(jìn)行比較。這種方法繪制了更多的區(qū)域，并給出了更豐富多彩的結(jié)果。

扼要重述

Unet是一個(gè)非常酷的模型。在第一個(gè)分割任務(wù)中，我們?cè)谟?xùn)練過程中遇到了一個(gè)問題，我們使用高分辨率圖像，這就是我們使用就地BatchNorm的原因。在我們的第二個(gè)任務(wù)（修復(fù)）中，我們使用了部分卷積而不是默認(rèn)卷積，這使我們能夠獲得更好的結(jié)果。在進(jìn)行彩色化時(shí)，我們添加了一個(gè)小的鑒別器網(wǎng)絡(luò)，該網(wǎng)絡(luò)對(duì)不真實(shí)圖像的生成器進(jìn)行懲罰。我們還使用了知覺損失。
第二個(gè)結(jié)論-評(píng)估師是必不可少的。不僅在創(chuàng)建分段掩碼階段，而且在最終結(jié)果驗(yàn)證階段。最后，我們給用戶三張圖片：一張有修補(bǔ)缺陷的原始圖片，一張有修補(bǔ)缺陷的彩色圖片和一張簡(jiǎn)單的彩色圖片，以防缺陷搜索和修補(bǔ)算法出錯(cuò)。
我們從戰(zhàn)爭(zhēng)相冊(cè)項(xiàng)目中拍攝了一些照片，并通過這些神經(jīng)元進(jìn)行處理。以下是我們得到的結(jié)果：

此外，在這里您可以更仔細(xì)地查看原始圖像和所有處理階段。