AI-Based Photo Restoration

大家好！我是mail.ru集團計算機視覺團隊的研究工程師。在這篇文章中，我將講述一個關于我們如何為舊的軍事照片創建基于人工智能的照片修復項目。什么是“照片復原”？它包括三個步驟：
?發現所有圖像的缺陷：破裂、刮傷、孔洞；
?根據周圍的像素值修復發現的缺陷；
?給圖像上色。
此外，我將描述照片恢復的每一步，并告訴你我們如何獲得數據，我們訓練了哪些網絡，我們完成了什么，以及我們犯了什么錯誤。

尋找缺陷

我們想在上傳的照片中找到所有與缺陷相關的像素。首先，我們需要弄清楚人們會上傳什么樣的圖片。我們采訪了“不朽軍團”項目的創始人，這是一個非商業組織，存儲了二戰時遺留的照片，他們與我們分享了他們的數據。經過分析，我們注意到人們上傳的大多是個人或團體肖像，存在中等到大量的缺陷。
然后我們得收集一套訓練器材。分割任務的訓練集是一個圖像和一個標記所有缺陷的掩碼。最簡單的方法是讓評估人員創建分段掩碼。當然，人們很清楚如何發現缺陷，但這需要很長時間。

在一張照片上標記缺陷像素可能需要一個小時或整個工作日。因此，要在幾周內收集到超過100張圖片的訓練集并不容易。這就是為什么我們試圖增加我們的數據并創建我們自己的缺陷：我們會拍一張好的照片，在圖像上使用隨機游動添加缺陷，最后得到一個顯示有缺陷的圖像部分的遮罩。在沒有增強的情況下，我們在訓練集中有68張手動標記的照片，在驗證集中有11張。

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

當我們在任務中使用這種分割方法時會出現什么問題？

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

網絡在18小時內匯聚在4個Geforce 1080ti上。推理需要290毫秒。這很長，但這是我們優于默認性能的代價。Validation DICE equals 0,35, and ROCAUC — 0,93.

圖像修復

與分割任務一樣，我們使用Unet。要進行修補，我們需要上傳一張原始圖像和一個遮罩，在那里我們用1和0標記所有清潔區域-所有我們想要修補的像素。這就是我們收集數據的方式：對于來自開源圖像數據集（例如openimagesv4）的任何照片，我們都添加了與現實生活中看到的類似的缺陷。然后我們訓練網絡恢復丟失的部分。
我們如何修改這個任務的Unet？
我們可以用部分卷積代替原來的卷積。其思想是，當我們用一些核卷積一個區域時，我們不考慮缺陷像素值,這使得修補更加精確。我們向您展示了最近Nvidia論文中的一個例子。他們在中間的圖片中使用了默認的二維卷積的Unet，在右邊的圖片中使用了部分卷積的Unet。

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

著色

我們將缺陷分割并修補；第三步-顏色重建。就像我之前說的，照片中有很多個人和團體的肖像。我們希望我們的網絡能和他們很好地配合。我們決定提出自己的彩色化，因為沒有一個現有的服務可以快速和有效的彩色肖像。我們希望彩色照片更真實可信。

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

為什么我們選擇YCbCr代表？人眼對亮度的變化比對顏色的變化更敏感。這就是為什么我們重復使用人眼最敏感的Y分量（亮度），并預測Cb和Cr，因為我們不能很好地注意到顏色的虛假性，所以我們可能會犯錯誤。這一特性在彩色電視出現之初就得到了廣泛的應用，當時的頻道容量還不足以傳輸所有的顏色。圖像在YCbCr中傳輸，不改變y分量，Cb和Cr減少一半。

如何創建基線

我們可以采用帶預訓練編碼器的Unet，并將現有CbCr值和預測值之間的l1損失最小化。我們想給人像上色，因此，除了openimages照片，我們還需要更多特定于任務的照片。
我們在哪里可以得到穿著軍裝的人的彩色照片？互聯網上有人把舊照片上色當作一種愛好或是一種代價。他們做得很仔細，盡量精確。當他們給制服、肩章和獎章上色時，他們會參考檔案材料，因此他們的工作成果是值得信賴的。總而言之，我們使用了200張手工著色的照片，上面有穿著軍裝的人。
另一個有用的數據來源是工農紅軍網站。它的一個創立者幾乎每到第二次世界大戰時都會穿著蘇維埃的制服拍照。

在一些照片中，他模仿著名的檔案照片中人們的姿勢。他的照片有白色背景是一件好事：它讓我們可以通過在背景中添加各種自然物體來很好地擴充數據。我們還使用了一些常規的肖像，以徽章和其他戰時特征作為補充。
我們訓練了Albunet-50——它是一個使用預先訓練的Resnet-50作為編碼器的Unet。網絡開始提供足夠的結果：皮膚是粉紅色的，眼睛是灰藍色的，肩膀是黃色的。然而，問題是它讓照片上的一些區域沒有受到影響。這是因為根據錯誤L1，在不做任何事情比試圖預測某種顏色更好的地方找到這樣的最佳值。

我們正在將我們的結果與地面真實照片進行比較-這是由Klimbim完成的手動著色。
我們怎樣才能解決這個問題？我們需要一個鑒別器：一個接收圖像并告訴我們圖像是否真實的神經網絡。下面的一張圖片是手動著色的，另一張是由我們的生成器Albunet-50著色的。人類如何區分手動和自動彩色照片？通過觀察細節。你能告訴我們的基線解決方案后自動著色的照片在哪張圖嗎？

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

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

扼要重述

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

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