Keras實現(xiàn)：用部分卷積補(bǔ)全圖像不規(guī)則缺損

編者按：圖像補(bǔ)全是一個熱門研究領(lǐng)域，今年4月，NVIDIA發(fā)表了一篇精彩的論文：Image Inpainting for Irregular Holes Using Partial Convolutions。文章指出，以往圖像補(bǔ)全技術(shù)都是用殘缺位置周邊的有效像素統(tǒng)計信息填充目標(biāo)區(qū)域，這種做法雖然結(jié)果平滑，但存在效果不逼真、有偽像，且后期處理代價昂貴的缺點。因此他們用大量不規(guī)則掩膜圖像訓(xùn)練了一個深度神經(jīng)網(wǎng)絡(luò)，它能為圖像生成合理掩膜，再結(jié)合僅以有效像素為條件的部分卷積（Partial Convolutions），最終模型的圖像補(bǔ)全效果遠(yuǎn)超前人的成果。

而近日，有網(wǎng)友復(fù)現(xiàn)了這篇論文，并在GitHub上公開了他的Keras實現(xiàn)，感興趣的讀者前去一看：github.com/MathiasGruber/PConv-Keras

NVIDIA論文

環(huán)境

Python 3.6

Keras 2.2.0

Tensorflow 1.8

如何使用這個repo

repo中包含的PConv2Dkeras實現(xiàn)可以在libs/pconv_layer.py和libs/pconv_model.py中找到。除此之外，作者還提供了四個jupyter NoteBook，詳細(xì)介紹了實現(xiàn)網(wǎng)絡(luò)時經(jīng)歷的幾個步驟，即：

step 1：創(chuàng)建隨機(jī)不規(guī)則掩膜

step 2：實現(xiàn)和測試PConv2D層

step 3：實現(xiàn)和測試采用UNet架構(gòu)的PConv2D層

step 4：在ImageNet上訓(xùn)練和測試最終模型

實現(xiàn)細(xì)節(jié)

在設(shè)計圖像補(bǔ)全算法時，研究人員首先要考慮兩個因素：從哪里找到可利用的信息；怎么評判整體補(bǔ)全效果。無論是天然破損的圖像，還是被人為打上馬賽克的圖像，這之中都涉及圖像語義上的預(yù)測。

這篇論文發(fā)表之前，學(xué)界在圖像補(bǔ)全上的最先進(jìn)方法之一是利用剩余圖像的像素統(tǒng)計信息來填充殘缺部分，這利用了同一幅圖像素間的連接性，但缺點是只反映了統(tǒng)計上的聯(lián)系，無法真正實現(xiàn)語義上的估計。后來也有人引入深度學(xué)習(xí)的方法，訓(xùn)練了一個深度神經(jīng)網(wǎng)絡(luò)，以端到端的方式學(xué)習(xí)語義先驗和有意義的隱藏表示，但它仍局限于初始值，而且使用的是固定的替換值，效果依然不佳。

NVIDIA在論文中提出了一種新技巧：添加部分卷積層（Partial Convolutional Layer），并在這一層之后加上一個掩膜更新步驟。部分卷積層包含生成掩膜和重新歸一化，它類似圖像語義分割任務(wù)中的segmentation-aware convolutional（分段感知卷積），能在不改變輸入掩膜的情況下分割圖像信息。

簡而言之，給定給定一個二元掩膜，部分卷積層的卷積結(jié)果只取決于每一層的非殘缺區(qū)域。相比segmentation-aware convolutional，NVIDIA的創(chuàng)新之處是自動掩膜更新步驟，它可以消除部分卷積能夠在非掩膜值上操作的任何掩膜。

具體設(shè)計過程可以閱讀論文查看，下面我們只總結(jié)一些細(xì)節(jié)。

生成掩膜

為了訓(xùn)練能生成不規(guī)則掩模的深度神經(jīng)網(wǎng)絡(luò)，論文研究人員截取視頻中的兩個連續(xù)幀，用遮擋/解除遮擋創(chuàng)建了大量不規(guī)則掩膜，雖然他們在論文中稱將公開這個數(shù)據(jù)集，但現(xiàn)在我們還找不到相關(guān)資源。

在這個Keras實現(xiàn)中，作者簡單創(chuàng)建了一個遮擋生成函數(shù)，用OpenCV繪制一些隨機(jī)的不規(guī)則形狀，以此作為掩膜數(shù)據(jù)，效果目前看來還不錯。

部分卷積層

這個實現(xiàn)中最關(guān)鍵的部分就是論文的重點“部分卷積層”?；旧?，給定卷積filter W和相應(yīng)的偏差b，部分卷積的形式是：

其中⊙表示點乘，即每個矩陣元素對應(yīng)相乘，M是由0和1構(gòu)成的二進(jìn)制掩碼。在每次完成部分卷積操作后，掩膜要進(jìn)行一輪更新。這意味著如果卷積能夠在至少一個有效輸入上調(diào)節(jié)其輸出，則在該位置移除掩碼：

這樣做的結(jié)果是，在網(wǎng)絡(luò)夠深的情況下，最終掩碼將全部為0（消失）。

UNet架構(gòu)

下圖是論文中提供的PConv整體架構(gòu)，它類似UNet，只不過其中所有正常的卷積層都被部分卷積層代替，使圖像+掩膜無論何時都能一起通過網(wǎng)絡(luò)

PConv整體架構(gòu)

PConv彩色圖解

損失函數(shù)

論文涉及的損失非常多樣，簡而言之，它包括：

掩膜區(qū)（Lhole）和非掩膜區(qū)（Lvalid）的每個像素?fù)p失

基于ImageNet預(yù)訓(xùn)練的VGG-16（pool1, pool2 and pool3 layers）的感知損失（Lperceptual）

VGG-16在預(yù)測圖像（Lstyleout）和計算圖像（Lstylecomp）上的風(fēng)格損失（以非殘缺區(qū)像素為真實值）

殘缺區(qū)域每個像素擴(kuò)張的總變差損失（Ltv），也就是1像素擴(kuò)張區(qū)域的平滑懲罰

以上損失的權(quán)重如下：

論文補(bǔ)全效果

下圖是論文中呈現(xiàn)的圖像補(bǔ)全效果，其中第一列是輸入圖像，第二列是GntIpt算法的輸出，第三列是NVIDIA論文的結(jié)果，第四列是真實完整圖像。可以發(fā)現(xiàn)，無論圖像缺失區(qū)域有多不規(guī)則，PConv的補(bǔ)全效果在顏色、紋理、形狀上都更逼真，也更平滑流暢。

小結(jié)

最后一點，也是最重要的一點，如果是在單個1080Ti上訓(xùn)練模型，batch size為4，模型訓(xùn)練總用時大約在10天左右，這是個符合論文所述的數(shù)據(jù)。所以如果有讀者想上手實踐，記得提前做好硬件和時間上的準(zhǔn)備。