Image Inpainting with Learnable Bidirectional Attention Maps
Abstract
大多數基于卷積網絡(CNN)的修複方法采用标準的卷積方法對有效像素和空洞進行不可區分的處理，使其僅限于處理不規則的空洞，更容易産生色差和模糊的修複結果。部分卷積被提出用來解決這一問題，但它采用手工特征重歸一化，并且隻考慮前向掩碼更新。本文提出了一種可學習的注意力圖子產品，用于端到端的特征重歸一化和掩碼更新的學習，能夠有效地适應不規則孔洞和卷積層的傳播。此外，引入了可學習的反向注意圖，使U-Net的解碼器專注于填充不規則的空洞，而不是同時重構空洞和已知區域，進而得到了可學習的雙向注意圖。定性和定量實驗表明，我們的方法在生成更清晰、更連貫和視覺上可信的修複結果方面表現得比最先進的方法更好。源代碼和預先教育訓練的模型可在以下網址獲得：https://github.com/Vious/LBAM_inpainting/。
1 Intorduction
圖像修複[3]是一種具有代表性的低層視覺任務，其目标是填充圖像中的空洞，具有許多現實應用，如分散目标的去除、遮擋區域的填充等。然而，對于圖像中給定的孔，可能存在多個可能的解，即可以用與周圍已知區域一緻的任何看似合理的假設來填充這些孔。而且孔洞可能是複雜和不規則的圖案，進一步增加了圖像修複的難度。傳統的基于範例的方法[2，18，32]，例如PatchMatch[2]，通過從已知區域搜尋并複制相似的面片來逐漸填充孔洞。盡管基于樣本的方法在生成細節紋理方面是有效的，但它們在捕獲進階語義方面仍然有限，并且可能無法生成複雜和非重複的結構(參見圖1©)。
近年來，深度卷積網絡(CNNs)在圖像修複中的應用取得了長足的進展[10，20]。基于CNN的方法得益于強大的表征能力和大規模的訓練，能夠有效地實作語義上合理結果的生成。對抗性損失[8]也被用來改善結果的感覺品質和自然性。盡管如此，大多數現有的基于CNN的方法通常采用标準卷積，無法區分有效像素和空洞。是以，它們在處理不規則孔方面受到限制，并且更有可能産生具有顔色差異和模糊的修複結果。作為補救措施，已經引入了幾種後處理技術[10，34]，但仍然不足以解決僞影(參見圖1(D))。基于CNN的方法也與基于樣本的方法相結合，以顯式地合并孔洞的掩模，以更好地恢複結構和增強細節[26，33，36]。在這些方法中，掩模被用來引導編碼器特征從已知區域傳播到孔。然而，複制和增強操作大大增加了計算成本，并且僅部署在一個編解碼層。是以，它們在填充矩形孔洞方面做得更好，而在處理不規則孔洞方面表現不佳(參見圖1(E))。
為了更好地處理不規則孔并抑制顔色差異和模糊，提出了部分卷積(PConv)[17]。在每個PConv層中，使用掩碼卷積使輸出僅以未掩碼的輸入為條件，并引入特征重新歸一化來縮放卷積輸出。進一步提出了一種掩碼更新規則來更新下一層的掩碼，使得PConv在處理不規則孔洞時非常有效。盡管如此，PConv采用了硬0-1掩碼和手工制作的特征重新歸一化通過絕對信任所有填充的中間特征。此外，部分卷積僅前向掩碼更新，并簡單地将所有掩碼用于解碼器特征。
在本文中，我們更進一步，提出了可學習的雙向注意力圖子產品，用于U-Net[22]體系結構的編碼器和解碼器上的特征的重新歸一化。首先，我們無偏見地重溫了PConv，并證明可以安全地避免掩碼卷積，并且特征重新歸一化可以解釋為硬0-1掩碼引導的重新歸一化。為了克服硬0-1掩碼和手工掩碼更新的局限性，我們提出了一種可學習的注意力圖子產品，用于學習特征重歸一化和模闆更新。通過端到端的訓練，可學習的注意力圖能有效地适應不規則的孔洞和卷積層的傳播。
此外，PConv在解碼器功能上簡單地使用了全一掩模，使得解碼器應該同時出現洞和已知區域的幻覺。注意，已知區域的編碼器特征将被連接配接起來，解碼器自然隻需要專注于孔的修複。是以，我們進一步引入了可學習的反向注意圖，使得U-Net的解碼器隻專注于填補空洞，進而得到了可學習的雙向注意圖。與PConv相比，經驗性地部署可學習的雙向注意力圖有利于網絡訓練，使得包含對抗性損失以提高結果的視覺品質是可行的。
在Paris SteetView[6]和Places[40]資料集上進行了定性和定量實驗，以評估我們提出的方法。結果表明，我們提出的方法在生成更清晰、更連貫和視覺上可信的修複結果方面表現良好。從圖1(F)(G)看，與PConv相比，我們的方法在生成清晰的語義結構和逼真的紋理方面更有效。綜上所述，這項工作的主要貢獻有三個方面：
a、提出了一種用于圖像修複的可學習注意力圖子產品。與PConv相比，可學習的注意力圖更能适應任意不規則的孔洞和卷積層的傳播。
b、正向注意力圖和反向注意力圖結合在一起構成了可學習的雙向注意力圖，進一步提高了結果的視覺品質。
c、在兩個資料集上的實驗和真實世界的物體去除實驗表明，我們的方法在生成整形、更連貫和視覺上的結果上比最先進的方法有更好的表現。
2 Related Work
在這一部分中，我們簡要介紹了相關工作，特别是基于樣本的方法所采用的傳播過程以及基于CNN的修複方法的網絡結構。
2.1 Exemplar-based Inpainting
大多數基于樣本的修複方法都是從已知區域搜尋并粘貼，從外到内逐漸填充孔洞[2，4，18，32]，其結果高度依賴于傳播過程。一般說來，通過先填充結構，再填充其他缺失區域可以獲得更好的修複效果。為了指導更新檔處理順序，引入了更新檔優先級[15，29]測度作為置信度和資料項的乘積。雖然置信項通常定義為輸入圖檔中已知像素的比率，但是已經提出了幾種形式的資料項。特别是Criminisi等人[4]提出了一種基于梯度的資料術語，用于填充優先級較高的線性結構。Xu和Sun[32]假設結構塊在圖像中是稀疏分布的，并提出了一種基于稀疏性的資料項。Le Meuret al等人[18]采用結構張量的特征值差異[5]作為結構更新檔的名額。
2.2 Deep CNN-based Inpainting
早期基于CNN的方法[14，21，30]被建議用于處理帶有小孔和細孔的圖像。在過去的幾年裡，深度神經網絡受到了人們的極大關注，并在填充大孔洞方面表現出了良好的性能。
Phatak等人[20]采用了編解碼器網絡(即上下文編碼器)，并考慮了重構和對抗性損失，以更好地恢複語義結構。Iizuka等人[10]結合了全局和局部鑒别器，既能再現語義上合理的結構，又能再現局部真實的細節。Wang等人。[28]提出了一種結合信心驅動的重建損失和隐式多樣化MRF(ID-MRF)項的生成式多列CNN。多階段方法也已被研究，以減輕訓練深度修複網絡的難度。張等人。[37]提出了一種漸進式産生式網絡(PGN)來填補多階段的空洞，而LSTM則用來開發階段間的依賴關系。Nazeri等人的研究成果。[19]提出了一種兩階段模型EdgeConnect，該模型首先預測顯著的邊緣，然後由邊緣引導生成修複結果。相反，熊等人[31]針對結構推理和内容生成的關系，提出了前景感覺修複，包括輪廓檢測、輪廓補全和圖像補全三個階段。為了将基于樣本的方法和基于CNN的方法結合起來，Yang 等人[34]提出了多尺度神經塊合成方法(MNPs)，通過結合整體内容和局部紋理限制的聯合優化來優化上下文編碼器的結果。進一步開發了其他兩階段前饋模型，例如上下文關注[26]和更新檔交換[36]，以克服MNP的高計算成本，同時顯式地利用已知區域的圖像特征。同時，嚴等人[33]将U-Net修改為單階段網絡，即shift-net，以利用編碼器特征從已知區域的移位來更好地再現似是而非的語義和詳細内容。最近，鄭等人[39]引入了增強的短期+長期注意力層，提出了一種具有兩條并行路徑的多元修複機率架構。
大多數現有的基于CNN的修複方法通常不太适合處理不規則的孔。為了解決這個問題，劉等人[17]提出了一種部分卷積(PConv)層，包括掩碼卷積、特征重歸一化和掩碼更新三個步驟。Yu等人[35]提供了門控卷積，其通過考慮損壞的圖像、掩模和使用者草圖來學習通道式軟掩碼。然而，PConv采用手工制作的特征重整化，隻考慮前向掩碼更新，在處理顔色差異和模糊方面仍然有限(見圖1(D))。
3 Proposed Method
在這一部分中，我們首先回顧PConv，然後介紹我們的可學的習雙向注意力圖。随後，給出了該方法的網絡結構和學習目标。
3.1 Revisiting Partial Convolution
PConv[17]層通常包括三個步驟，即(i)掩碼卷積、(ii)特征重新歸一化和(iii)掩碼更新。進一步設W為卷積濾波，b為其偏置。首先，我們引入卷積掩模Mc=M⊗k1/9，其中⊗表示卷積算子，k1/9表示每個元素1/9的3×3卷積濾波器。PConv的過程可以表示為：
其中，A=FA(Mc)表示注意圖，M’=fM(Mc)表示更新的掩碼。我們進一步将注意力圖和更新掩碼的激活函數定義為：
來自等式(1)∼(5)和圖2(A)中，PConv也可以解釋為掩碼和圖2(A)之間的一種特殊的互相作用模型和卷積特征圖。然而，PConv采用了手工制作的卷積濾波器k1/9以及手工制作的激活函數FA(Mc)和Fm(Mc)，進而為進一步改進提供了一些回旋餘地。此外，fM(Mc)的不可微特性也增加了端到端學習的難度。據我們所知，将對抗性損失與PConv結合訓練U-Net仍然是一個困難的問題。此外，PConv隻考慮掩碼及其對編碼器特性的更新。在解碼器功能上，簡單的采用全一掩碼，使得PConv僅限于填充空洞。
3.2 Learnable Attention Maps
無偏置卷積層在U網中已被廣泛采用，用于圖像到圖像的轉換[11]和圖像修複[33]。當偏置被去除時，它可以很容易地從方程(2)中看到，更新孔中的卷積特征為零。是以，等式(1)中的掩碼卷積等效地重寫為标準卷積:
然後，在等式(2)中對特征進行重新歸一化可以解釋為卷積特征和注意力圖的元素乘積：
即使手工制作的卷積濾波器k1/9是固定的，并且不适合該掩模。更新掩模的激活函數絕對信任區域Mc>0中的修複結果，但更明智的做法是，對Mc較高的地區賦予更高的信心。為了克服上述局限性，我們從三個方面提出了可學習注意力圖，它從三個方面對PConv進行了無偏見的概括。首先，為了使掩模能夠适應不規則孔洞和分層傳播，我們用分層的、可學習的卷積濾波器Km代替了k1/9。其次，代替硬0-1掩碼更新，我們将更新掩碼的激活函數修改為:
其中α≥0是一個超參數，我們将α設定為0.8。可以看出，當α=0時，GM(Mc)退化為Fm(Mc)。第三，我們引入了一種非對稱高斯變形形式作為注意力圖的激活函數：
在a、µ、γl和γr缺少可學習參數的情況下，我們将它們初始化為a=1.1，µ=2.0，γl=1.0，γr=1.0，并以端到端的方式學習它們。綜上所述，可學習注意圖采用等式。(6)步驟(I)，下兩步公式如下：
圖2(B)說明了可學習注意圖的互相作用模型。與PConv相比，我們的可學習注意力圖更加靈活，可以端到端地訓練，使其能夠有效地适應不規則的孔洞和卷積層的傳播。
3.3 Learnable Bidirectional Attention Maps
當将PConv與U-Net結合用于修複時，該方法[17]僅更新掩模以及用于編碼器特征的卷積層。然而，解碼器功能通常采用全一掩碼。是以，應該使用第(l+1)層編碼器特征和第(L−l−1)層解碼器特征兩者來幻覺已知區域和孔中的第(L−l)層解碼器特征。實際上，第l層編碼器特征将與第(L−l)層譯碼特征串聯在一起，我們隻能專注于孔中第(L−l)層譯碼特征的生成。
我們進一步将可學習的反向注意力圖引入到解碼器的特征中。用Mc e表示編碼器特征Fin e的卷積掩碼。設Mc d=Md⊗kMd是解碼器特征Fin d的卷積掩碼。可學習反向注意圖的前兩個步驟可以表示為：
其中，We和Wd是卷積濾波器。我們将Ga(Mcd)定義為反向注意力圖。然後，更新掩碼Mc_Di并将其部署到先前的解碼層：
圖2©說明了反向注意力圖的互相作用模型。與前向注意力圖不同，它同時考慮了編碼器特征(掩碼)和解碼器特征(掩碼)。此外，反向注意力圖中的更新掩碼被應用于前一解碼器層，而正向注意力圖中的更新掩碼被應用于下一編碼層。
通過将正向和反向注意力圖與U-net相結合，圖3顯示了完全可學習的雙向注意力圖。給定一幅帶有不規則孔洞的輸入圖像，我們使用Min表示二進制掩模，其中1表示有效像素，0表示孔洞中的像素。從圖3開始，前向注意力圖将Min作為第一層編碼器特征重新歸一化的輸入掩碼，并逐漸更新并應用到下一編碼層。相反，反向注意力圖以1−Min作為輸入，對解碼特征的最後一層(即第L層)進行重新歸一化，并逐漸更新并将掩碼應用到前一層解碼器層。得益于端到端的學習，我們的可學習雙向注意圖(LBAM)在處理不規則洞時更有效。反向注意力圖的引入使得解碼器隻專注于填充不規則的空洞，這也有助于修複性能。我們的LBAM還有利于網絡訓練，使得利用對抗性損失來提高視覺品質是可行的。
3.4 Model Architecture
我們通過去除瓶頸層并結合雙向注意圖來修改14層的U-Net體系結構11。特别地，前向注意力層應用于編碼器的前六層，而反向注意力層應用于後六層解碼器。對于所有的U-Net層和正向和反向關注層，我們使用了核大小為4×4、步長為2、填充為1的卷積濾波器，并且不使用偏置參數。在U-Net主幹網中，對重新歸一化後的特征進行批歸一化和leaky RELU非線性處理，在最後一層卷積後立即展開tanh非線性。
3.5 Loss Functions
為了更好地恢複紋理細節和語義，我們綜合了像素重構損失、感覺損失[12]、樣式損失[7]和對抗性損失[8]來訓練LBAM。
Pixel Reconstruction Loss.用Iin表示帶孔的輸入圖像，用Mini表示二進制掩模區域，用Ig表示地面真實圖像。我們的LBAM的輸出可以定義為IOUT=Φ(Iin，Min；Θ)，其中Θ表示要學習的模型參數。我們采用輸出圖像的ℓ1範數誤差作為像素重建損失:
Perceptual Loss. ℓ1範數損失僅限于捕捉進階語義，與人類對圖像品質的感覺不一緻。為了緩解這一問題，我們引入了在ImageNet[23]上預先訓練的VGG-16網絡[25]上定義的感覺損失Lpert:
其中PI(·)是第i個池層的特征地圖。在我們的實施中，我們使用預先教育訓練的VGG-16的Pool-1、Pool-2和Pool-3層。
Style Loss.為了更好地恢複細節紋理，我們進一步采用了VGG-16彙聚層特征圖上定義的樣式損失。類似于[17]，我們從特征地圖的每一層構造一個Gram矩陣。假設特征圖PI(I)的大小為Hi×Wi×Ci。然後，可以将樣式損失定義為:
Adversarial Loss.對抗性損失[8]已廣泛應用于圖像生成[24，27，38]和底層視覺[16]中，以提高生成圖像的視覺品質。為了提高GaN的訓練穩定性，Arjovskite等人[1]利用Wasserstein距離來度量生成圖像和真實圖像之間的分布差異，以及Gulrajani等人[9]進一步引入梯度懲罰來加強鑒别器中的Lipschitz限制。在[9]之後，我們将對抗性損失表示為：
其中D(·)表示鑒别器。I是通過随機選擇因子的線性插值從Ig和Iout中獲得的，在我們的實驗中，λ被設定為10。我們經驗發現，當包含對抗性損失時，很難訓練PConv模型。幸運的是，結合可學習的注意力圖有助于訓練，使得在對抗性損失下學習LBAM成為可能。
Model Objective. 考慮到上述損失函數，我們的LBAM的模型目标可以表示為:
其中，λ1、λ2、λ3和λ4是權衡的參數。在我們的實作中，我們經驗地設定了λ1=1、λ2=0.1、λ3=0.0 5和λ4=120。
4. Experiments
在巴黎街景[6]和Places(Places365-Standard)[40]這兩個資料集上對我們的LBAM算法進行了實驗，這兩個資料集在圖像修複文獻[20，33，34，36]中得到了廣泛的應用。對于巴黎街景，我們使用其原始拆分，14,900張圖像用于教育訓練，100張圖像用于測試。在我們的實驗中，随機選擇100幅圖像并從訓練集中移除，以形成我們的驗證集。至于地點，我們從365個類别中随機選擇10個類别，并使用原始訓練集中每個類别的所有5000張圖像來構成我們的5萬張圖像的訓練集。此外，我們将每類1000幅圖像中的原始驗證集分成兩個相等且互不重疊的500幅圖像集，分别用于驗證和測試。我們的LBAM處理256×256圖像所需的∼時間為70ms，速度提高了5倍比上下文關注[36]和∼快3倍于全局和局部(GL)[10]。
在我們的實驗中，所有的圖像都被調整了大小，其中最小的高度或寬度是350，然後随機裁剪到256×256的大小。訓練過程中采用翻轉等資料增強。我們生成了18,000個随機形狀的面具，并從[17]中生成了12,000個面具用于訓練和測試。我們的模型使用Adma算法[13]進行了優化，初始學習率為1E−4，β=0.5。訓練過程在500個曆元後結束，并且最小批大小為48。所有實驗都是在配備4個并行NVIDIA GTX 1080Ti GPU的PC上進行的。
4.1 Comparison with State-of-the-arts
我們的LBAM與四種最先進的方法進行了比較，即全局和局部[10]、PatchMatch[2]、上下文注意[36]和PConv[17]。
Evaluation on Paris StreetView and Places. 圖4和圖5顯示了我們的LBAM和競争方法的結果。GLOBAL&LOCAL[10]在處理不規則孔方面受到限制，會産生許不比對比且毫無意義的紋理。PatchMatch[2]在恢複複雜結構方面表現不佳，結果與周圍環境不一緻。對于一些複雜和不規則的孔，上下文注意[36]仍然會産生模糊的結果，并且可能産生不需要的僞影。PConv[17]在處理不規則孔時是有效的，但在某些區域仍然不可避免地會出現過度平滑的結果。相比之下，我們的LBAM表現很好，生成了視覺上更合理的結果，具有精細的細節和逼真的紋理。
Quantitative Evaluation. 我們還在掩模比為(0.1，0.2)，(0.2，0.3)，(0.3，0.4)和(0.4，0.5)的位置[40]上與競争的LBAM方法進行了定量比較。從表1可以看出，我們的LBAM在PSNR、SSIM和平均ℓ1損耗方面表現良好，特别是當掩模比大于0.3時。
Object Removal from Real-world Images. 使用基于位置訓練的模型，我們進一步評估了LBAM在現實世界目标移除任務中的性能。圖6顯示了我們的LBAM、上下文注意[36]和PConv[17]的結果。我們使用輪廓形狀或矩形邊界框來遮罩對象區域。與同類方法相比，我們的LBAM能夠同時利用全局語義和局部紋理産生逼真和連貫的内容。
User Study. 此外，還對巴黎街景和主觀視覺品質評價場所進行了使用者研究。我們從覆寫着不同不規則孔洞的測試集中随機選取了30幅圖像，修複結果由PatchMatch[2]、Global&Local[10]、Context Attendence[36]、PConv[17]和我們的算法生成。我們邀請了33名志願者投票選出視覺上最合理的修複結果，評價标準包括與周圍環境的一緻性、語義結構和精細性細節。
對于每個測試圖像，5個修複結果被随機排列并與輸入圖像一起呈現給使用者。我們的LBAM有63.2%的機會赢得最有利的結果，大大超過了PConv17，PatchMatch2，上下文關注度36和Global&Local10。
4.2 Ablation Studies
消融研究是為了比較幾種LBAM變體在巴黎街景上的性能，即(I)我們的(完整)：完整的LBAM模型，(ii)我們的(未學習)：LBAM模型，其中掩碼卷積濾波器的所有元素都設為1/16，因為濾波器的大小為4×4，并且我們采用公式中定義的激活函數(4)和Eqn.(5)，(iii)我們的(前向)：沒有反向注意圖的LBAM模型，(iv)我們的(w/o Ladv)：無(w/o)對手損失的LBAM模型，V)我們的(Sigmoid/LReLU/RELU/3×3):LBAM模型使用Sigmoid/LeakyReLU/REU作為激活函數或3×3濾波器進行掩碼更新。
圖7顯示了我們的(未學習的)來自第一編碼器層和第13解碼器層的特征的可視化，我們的(前向更新)和我們的(全部)。對于我們的(未學到的)，模糊和僞影可以從圖9(B)中觀察到。我們的(前向更新)有利于減少僞影和噪聲，但解碼器會同時對孔和已知區域産生幻覺，并産生一些模糊效果(參見圖9©)。相比之下，我們的(Full)在生成語義結構和詳細紋理方面是有效的(見圖9(D))，解碼器主要關注生成孔(見圖7(G))。表2給出了LBAM變體在巴黎街景上的定量結果，我們的性能增益(FULL)可以用(1)可學習的注意力圖，(2)反向注意力圖，和(3)适當的激活函數來解釋。
Mask Updating. 圖8顯示了來自不同層的更新遮罩的可視化。從第一層到第三層，編碼器的掩模逐漸更新，以減小孔洞的大小。類似地，從第13層到第11層，解碼器的掩碼逐漸更新，以減小已知區域的大小。
Effect of Adversarial Loss. 表2還給出了有/無Ladv的定量結果。雖然我們的(無Ladv)改善了PSNR和SSIM，但Ladv的使用通常有利于修複結果的視覺品質。定性的結果在附錄中給出。
5. Conclusion
提出了一種用于圖像修複的可學習雙向注意圖(LBAM)。通過引入可學習的注意力圖，我們的LBAM能夠有效地适應不規則的孔洞和卷積層的傳播。此外，還提出了反向注意力圖，使U-Net的解碼器隻專注于空洞的填充。實驗表明，我們的LBAM在生成更清晰、更連貫和更精細的結果方面表現優于最先進的水準。