基于掩碼預測和多尺度上下文聚合的人臉圖像修復算法

孫劍明，吳金鵬，沈子成，彭俄禎

(哈爾濱商業(yè)大學 計算機與信息工程學院，黑龍江 哈爾濱 150028)

0 引言

圖像修復[1]是指利用已知上下文信息填充圖像中缺失的區(qū)域。盡管研究人員已經提出許多圖像修復方法,無論是傳統(tǒng)算法還是現(xiàn)有深度學習算法,在修復圖像時都難以同時恢復合理的內容和清晰的紋理。

早期的圖像修復工作[1-13]試圖通過基于擴散或基于補丁的算法來解決這一問題?；跀U散的算法將上下文信息從邊界傳播到等照度線方向的孔洞(缺失區(qū)域)中?；谘a丁的算法通過從未損壞的圖像區(qū)域或外部數(shù)據(jù)庫復制相似的補丁來合成缺失區(qū)域?，F(xiàn)有的深度模型在低分辨率圖像的語義修復方面已經表現(xiàn)出良好的效果。但在更高分辨率(如512 pixel×512 pixel)使用這些深度模型往往會生成不合理的內容且紋理細節(jié)不真實,阻礙了用戶在高分辨率圖像中的實際應用。

為了克服以上2種問題,提出了一種基于掩碼預測和多尺度上下文聚合的人臉圖像修復模型。它包括一個以編碼器、多尺度上下文聚合模塊(Multi-Scale Context Aggregation Module,MSCAM)、解碼器為框架的生成器網絡和一個以掩碼預測為訓練任務的判別器網絡。

為了在破損圖像生成合理的內容,通過堆疊多個MSCAM。通過使用各種擴張率的空洞卷積來利用遙遠的距離上下文信息進行上下文推理,同時通過拼接多個特征圖來實現(xiàn)多尺度的特征融合。使用跳躍連接將編碼器中每一層卷積層的輸出與解碼器對應位置的輸入在通道維度上拼接,使得圖像的上下文信息向更高層分辨率傳播,從而為缺失區(qū)域生成更合理的內容。

為了生成真實的紋理細節(jié),使用了一種新穎的掩碼預測任務來訓練判別器。大多數(shù)現(xiàn)有的深度模型使用了帶有譜歸一化的PatchGAN判別器,從而迫使判別器將修復圖像中的所有補丁塊預測為假,而忽略了那些缺失區(qū)域之外的補丁塊確實來自真實圖像的事實。因此,這些深度模型可能難以生成逼真的細粒度紋理。為了克服以上問題,使用掩碼預測的PatchGAN(Mask Prediction-PatchGAN,MP-PatchGAN)判別器,迫使判別器區(qū)分真實和生成的小塊(缺失區(qū)域)的紋理細節(jié)。換句話說,對于修復圖像,判別器期望從真實圖像中分割出合成的圖像塊。這樣的學習目標導致了一個更強的判別器,并且反過來促進生成器來合成逼真的細粒度紋理。

本文的貢獻如下:

① 首次提出了MSCAM和MP-PatchGAN判別器。MSCAM進行特征提取,融合了來自不同感受野的特征圖,可以用來捕捉遙遠距離的上下文信息和感興趣的模式增強上下文推理。MP-PatchGAN判別器迫使判別器區(qū)分真實和生成的小塊(缺失區(qū)域)的紋理細節(jié),進一步提升人臉面部修復性能。

② 設計了一種基于掩碼預測和多尺度上下文聚合的人臉圖像修復網絡,以便生成合理的內容和逼真的紋理。

③ 本文設計的模型與其他主流的基于深度學習模型進行定量、定性評估,且評價結果優(yōu)于其他深度模型。

1 相關工作

由于圖像修復對圖像編輯應用(如物體移除和圖像復原)的重要實用價值,圖像修復已然成為近幾十年來活躍的研究課題之一?，F(xiàn)有的修復算法可以分為2類:基于傳統(tǒng)的算法和基于深度學習的算法。

1.1 基于傳統(tǒng)的圖像修復算法

基于擴散的算法[1-6]沿著等照度線方向將上下文像素從邊界傳播到孔洞。具體地,在像素傳播期間,通過使用偏微分方程來施加許多邊界條件。然而,這些方法通常會引入擴散相關的模糊,因此無法完成大面積缺失區(qū)域。

基于補丁的算法[7-13]通常通過從已知的圖像上下文或外部數(shù)據(jù)庫復制和粘貼相似的補丁塊來合成丟失的內容。然而,這些方法在完成復雜場景的大面積缺失區(qū)域的語義修復方面存在不足。這是因為基于補丁的方法嚴重依賴于通過低級特征的逐片匹配。這種技術不能合成已知區(qū)域中不存在類似補丁的圖像。

傳統(tǒng)算法往往在修復面積較小、紋理結構較為簡單時有比較好的效果。一旦缺失區(qū)域比較大(30%以上)時,修復效果往往會特別差。這是因為傳統(tǒng)的圖像修復算法往往不能深層地理解圖像中的語義信息。因此,隨著2015年之后深度學習的火熱,越來越多從事圖像修復的研究人員開始采用深度學習方法以獲取更深層次的語義信息理解以及更高質量的修復效果。

1.2 基于深度學習的圖像修復算法

深度特征學習和對抗訓練的出現(xiàn)使得圖像修復取得了重大進展。與基于傳統(tǒng)算法相比,深度修復模型能夠為復雜的場景生成合理的內容和逼真的細粒度紋理。

Pathak等[14]提出了上下文編碼器(Context Encoder,CE)模型,利用潛在特征空間中的通道等寬全連接層(Channel-wise Fully Connected Layer)將編碼特征和解碼特征連接。該模型首次使用生成對抗網絡(Generative Adversarial Network,GAN)框架,生成器由編碼器、解碼器組成。取AlexNet前5層作為編碼器,解碼器部分則由反卷積實現(xiàn)從高維特征向圖像真實大小的轉變。通過CE模型已經能使街景、人臉生成有希望的結果。然而,由于使用了通道等寬全連接層,使得該模型只能處理固定大小(128 pixel× 128 pixel)的圖像。CE結構如圖1所示。

針對上述問題,Iizuka等[15]提出在全卷積網絡(Fully Convolutional Network,FCN)建立模型,使得該模型能處理任意大小的圖像。為了生成逼真的細粒度紋理,通過采用GAN的框架進行語義修復,已經取得了重大進展。通過生成器和判別器之間的博弈論最小最大博弈,基于GAN的修復模型能夠生成更清晰的紋理。為了進一步改進判別器網絡,通過全局和局部判別器進行聯(lián)合訓練,從而達到全局和局部一致性。FCN結構如圖2所示。

然而,由于FCN中判別器網絡使用了全連接層,局部判別器只能處理固定形狀的缺失區(qū)域。為了解決這一問題,Yu等[16]繼承了PatchGAN判別器, PatchGAN判別器旨在區(qū)分真實圖像的Patch和修復圖像的Patch。但基于PatchGAN的模型通常忽略了這樣一個事實:那些缺失區(qū)域之外的補丁塊確實來自真實圖像的事實,并且盲目地促進判別器來區(qū)分這些相同的補丁塊是假的,因此會削弱生成器生成缺失區(qū)域之外的真實塊的真實內容。除此之外,為了捕獲基于FCN的遠距離上下文,提出了上下文注意力模塊,以通過逐片匹配從上下文中找到感興趣的補丁塊。Yu等提出的網絡結構如圖3所示。

圖3 上下文注意力機制網絡結構Fig.3 Contextual attention mechanism network structure

Liu等[17]提出部分卷積層代替標準卷積來處理缺失區(qū)域內外顏色差異和偽影的問題。部分卷積將像素分為有效和無效像素,且只對有效像素做卷積。Yu等[18]在上下文注意力模塊的基礎上使用了門控卷積學習掩碼的分布,進一步提升了修復性能。Yi等[19]提出了上下文殘差聚合(Contextual Residual Aggregation,CRA)機制,該機制可以通過對上下文信息中的殘差進行加權聚合來生成缺失區(qū)域的高頻殘差,提出的CRA網絡僅需輸入低分辨率圖像,再將低分辨率修復結果和高頻殘差合并得到高分辨率圖像。

傳統(tǒng)的圖像修復算法只適用于修復與背景相似的紋理,但人臉圖像面部成分之間存在緊密的相關性,人臉修復結果應當具有全局語義的合理性,如眉毛對齊、眼睛對齊等。因此,傳統(tǒng)的圖像修復算法不適用于人臉圖像修復。

隨著GAN的發(fā)展,有關深度學習的人臉圖像修復算法性能不斷提升。現(xiàn)有基于深度學習的人臉圖像修復算法主要分為4類:基于全連接層的算法、基于FCN的算法、基于注意力機制的算法和基于掩碼更新的算法。文獻[14]提出的CE網絡使用通道等寬全連接層完成臉部特征的長程遷移,但由于使用了全連接層,其網絡只能使用固定大小的人臉圖像。針對這一問題,文獻[15]首次提出使用FCN進行人臉修復,使其能處理任意大小的人臉圖像。文獻[16]提出了上下文注意力模塊,該模塊具有優(yōu)秀的臉部特征長程遷移能力。但是,這個模塊大尺寸圖像會占用巨大的顯存,因此無法在多尺度特征上使用,修復結果仍然會出現(xiàn)全局語義信息不合理的情況。文獻[17-18]和文獻[20]使用不同掩碼更新策略進一步提升了局部語義修復性能。但是,其在全局語義上存在如左右眼不對稱等不合理的問題。

2 網絡結構

2.1 MSCAM

受Inception v3[21]網絡的啟發(fā),MSCAM將標準卷積層拆分為4個子層,每個子層擁有較少的輸出通道,如圖4所示。首先,使用4個不同padding大小(1、2、4、8)的ReflectionPad2d轉化為4個特征圖。其次,將具有256個輸出通道的標準卷積層拆分為4個子層,使得每個子層具有64個輸出通道。每個子層通過使用不同擴張率(1、2、4、8)的空洞卷積來執(zhí)行上一ReflectionPad2d層輸出特征圖的不同變換。具體來說,通過在連續(xù)位置之間引入零來擴展卷積核。使用較大的擴張率使得卷積核能夠“看到”輸入圖像的較大區(qū)域,而使用較小擴張率只能使卷積核關注于較小感受野的局部區(qū)域。隨后通過在通道維度上拼接多尺度感受野的不同變換。最后使用標準卷積層進行多尺度上下文聚合。通過MSCAM能夠聚合遙遠距離的上下文信息和感興趣的模式來增強上下文推理。

2.2 MP-PatchGAN判別器

在早期深度修復模型的判別器中,通常采用softmax函數(shù)輸出一個整張圖像是否為真實的概率。最近的深度修復模型使用了帶有譜歸一化的PatchGAN判別器。然而,使用softmax函數(shù)的GAN判別器對于不規(guī)則破損圖像修復是不適用的。盡管,PatchGAN判別器適用不規(guī)則破損圖像修復,但PatchGAN判別器會將修復圖像中的所有補丁塊預測為假,而忽略了那些缺失區(qū)域之外的補丁塊確實來自真實圖像的事實。因此,無論是softmax函數(shù)的PatchGAN判別器,還是PatchGAN判別器都可能導致邊緣處結構扭曲、產生偽影。為了促進生成器生成細粒度的紋理,本文算法使用MP-PatchGAN判別器。圖5為訓練不同判別器的說明。

圖5 訓練不同判別器的說明Fig.5 Description of different training discriminators

具體來說,對輸入掩碼進行下采樣,作為真實圖像掩碼預測任務的真實值。使用補丁級別的軟掩碼作為生成圖像掩碼預測任務的真實值,它通過高斯濾波獲得。通過計算判別器生成的真實圖像的預測掩碼和輸入掩碼均方誤差(Mean Squared Error,MSE),再通過計算判別器生成的修復圖像的預測掩碼和補丁級軟掩碼的MSE。將判別器的對抗性損失表示為:

Ex～Pdata[(D(x)-1)2],

(1)

式中:D為判別器,G為生成器,σ為下采樣和高斯濾波的合成函數(shù),z為修復圖像,m為掩碼,1為與掩碼相同大小的元素全為1的矩陣,x為真實圖像,σ(1-m)為補丁級別軟掩碼。相應地,生成器的對抗性損失表示為:

(2)

式中:D為判別器,G為生成器,z為修復圖像,m為掩碼,⊙為逐像素相乘。僅對缺失區(qū)域的合成塊的預測被用于優(yōu)化生成器。通過這樣的優(yōu)化,使得判別器從缺失區(qū)域之外的真實上下文中分割缺失區(qū)域的合成塊,從而增強判別器的性能,反過來可以幫助生成器合成更真實的紋理。

2.3 算法整體框架

基于掩碼預測和多尺度上下文聚合的人臉圖像修復算法在設計上采用了GAN框架,用編碼器-MSCAM-解碼器為框架構造了生成器。首先,通過將真實圖像和不規(guī)則掩碼圖像合并成受損圖像作為輸入圖像。然后,編碼器提取圖像高層次的語義信息。再通過堆疊8層MSCAM捕獲遙遠距離的上下文特征和豐富的感興趣模式來增強上下文推理。使用跳躍連接將編碼器每一卷積層的輸出和解碼器對應位置上的輸入在通道維度上做拼接,使得圖像的上下文信息向更高層分辨率特征圖傳播。最后通過解碼器使得提取的特征圖向真實圖像的分辨率大小轉變,通過重建損失、對抗損失、感知損失、風格損失的聯(lián)合損失函數(shù)訓練生成器模型。本文算法在判別器上改用MP-PatchGAN判別器,通過對抗損失進行判別器訓練。該判別器用于區(qū)分真實和修復的Patch,從而生成更為清晰的紋理。本文算法結構如圖6所示。

圖6 基于掩碼預測和多尺度上下文聚合的人臉圖像修復模型Fig.6 Face image inpainting model based on mask prediction and multi-scale context aggregation

3 損失函數(shù)

圖像修復優(yōu)化目標既要保證像素重建精度,又要保證修復圖像的視覺效果。為此,遵循大多數(shù)現(xiàn)有深度修復模型,選擇4個優(yōu)化目標:重建損失、對抗損失[22]、感知損失[23]和風格損失[24]。采用對抗損失來優(yōu)化判別器,采用4個損失的加權聯(lián)合損失來優(yōu)化生成器。

3.1 重建損失

首先,使用L1損失,以確保像素級別的重建精度,重建損失表示為:

(3)

3.2 感知損失和風格損失

由于感知損失和風格損失對于圖像修復的有效性已經得到廣泛驗證[25],將它們包括在內以提高感知重建的準確性。具體而言,感知損失旨在最小化修復圖像和真實圖像的激活圖之間的L1距離,感知損失表示為:

(4)

(5)

在本文算法優(yōu)化中,感知損失計算VGG19網絡前5層修復圖像和真實圖像的激活圖之間的L1距離。風格損失計算VGG-19網絡前4層修復圖像和真實圖像的深度特征的格拉姆矩陣之間的L1距離。

3.3 聯(lián)合損失

使用重建損失、對抗損失、感知損失和風格損失加權聯(lián)合損失函數(shù)優(yōu)化生成器模型。聯(lián)合損失函數(shù)表示為:

(6)

式中:λ1=1,λ2=0.01,λ3=0.1,λ4=250。

4 實驗及分析

4.1 數(shù)據(jù)集來源

本文在公開數(shù)據(jù)集CelebA-HQ中訓練和測試模型,CelebA-HQ包括30 000張分辨率為512 pixel×512 pixel的人臉圖像,其中訓練集26 000張,驗證集2 000張,測試集2 000張。本文使用了Liu等[17]提供的任意形狀掩碼數(shù)據(jù)集,其中包括12 000張不同破損比例的任意形狀掩碼圖像。數(shù)據(jù)集樣例如 圖7所示。

圖7 數(shù)據(jù)集樣例Fig.7 Samples of dataset

4.2 參數(shù)設置

本文所有實驗均在CentOS平臺下進行,本文算法基于PyTorch 1.12、CUDNN 8.6和CUDA 11.6實現(xiàn)。CPU為Intel Xeon Platinum 8375C 2.90 GHz,GPU為NVIDIA GeForce RTX 3090。使用Adam算法[26]優(yōu)化模型,動量衰減指數(shù)β1=0.5、β2=0.999,學習率為0.001,批量大小為16。

4.3 現(xiàn)有深度模型

深度模型及其縮寫和簡要介紹如下:

① CA[16]是一個由粗到細的兩階段模型。它使用一個基于Patch的非局部模塊,即上下文注意力模塊。

② PConv[17]采用提出部分卷積層代替標準卷積來處理缺失區(qū)域內外顏色差異的問題。

③ GatedConv[18]在上下文注意力模塊的基礎上,結合了用于圖像修復的門控卷積和SN-PatchGAN判別器。

④ CRA[19]采用提出上下文殘差聚合模塊,該模塊通過加權聚合上下文的殘差來產生缺失內容的高頻殘差。將高頻殘差和上采樣后生成圖像聚合來生成高分辨率圖像,且用同一組注意力分數(shù)多次注意力轉移。注意力分數(shù)的共享實現(xiàn)了更少的參數(shù)以及在模型訓練方面更好的效率。

⑤ MAT[20]是首個基于Transfomer的人臉修復網絡,其提出的多頭上下文注意力模塊能夠獲取長距離的上下文信息提升臉部修復能力。MAT是目前人臉修復領域里SOTA算法。

4.4 定性分析

為了證明本文提出的人臉修復算法的優(yōu)越性,將本文算法與CA、PConv、GatedConv、CRA和MAT五種算法的修復結果進行定性比較,如圖8所示。

圖8 不同深度模型在CelebA-HQ測試集上的修復結果Fig.8 Inpainting results of different depth models on CelebA-HQ test set

結果表明,大多數(shù)深度模型在完成強語義的極大缺失區(qū)域修復時效果不理想。具體來說,CA模型通過提出的上下文注意力模塊可以從已知圖像內容中尋找與待修復區(qū)域相似度最高的Patch,然后使用這個Patch的特征做反卷積從而重建該Patch,但如果已知圖像內容和待修復區(qū)域差距很大,往往會產生不合理的內容,人臉的五官會出現(xiàn)扭曲(如圖8(c)所示)。PConv模型通過提出的部分卷積對每一卷積層的Mask更新,從而優(yōu)化修復結果,將已知區(qū)域內的像素視為有效像素,將缺失區(qū)域內的像素視為無效像素,且只對有效像素進行卷積,使用部分卷積代替標準卷積,能在一定程度上避免缺失區(qū)域內外顏色差異的問題,但眼睛會產生偽影(如圖8(d)所示)。GatedConv模型通過提出的門控卷積優(yōu)化部分卷積中Mask更新機制,該Mask更新機制是可學習的,但GatedConv模型使用PatchGAN判別器驅使已知區(qū)域像素改變,進而影響修復結果(如圖8(e)所示)。CRA模型通過多次的注意力轉移導致了嚴重的偽影(如圖8(f)所示)。MAT模型通過多頭上下文注意力模塊以及Transfomer架構的優(yōu)勢,其獲取長距離的上下文信息和全局感受野也能帶來不俗的修復性能(如圖8(h)所示)。而本文通過提出的MSCAM來捕捉遙遠距離的上下文信息和更多感興趣的模式和訓練掩碼預測的判別器來增強判別器的性能,以此生成更合理的內容和更清晰的紋理(如圖8(g)所示)。

4.5 定量分析

為了客觀評價本文模型在人臉圖像修復效果,對CA、PConv、GatedConv、CRA、MAT和本文提出的模型進行定量比較。評價指標包括L1Loss、峰值信噪比(Peak Signal to Noise Ratio,PSNR)、結構相似性(Structural Similarity,SSIM)和弗雷歇初始距離(Fréchet Inception Distance,FID)(如式(7)、式(8)、式(10)、式(11)所示),分別衡量修復圖像和原始圖像像素級差異、整體相似度、感知相似度以及特征相似度,其中L1Loss、FID 越低修復效果越好,PSNR、SSIM越高修復效果越好。具體來說,使用了CelebA-HQ驗證集的所有圖像。在不同孔洞率的掩碼下,對上述方法分別做測試。對于每個測試圖像,隨機采用自由形狀的掩碼作為測試掩碼。為了公平起見,在掩碼孔洞率一樣的情況下,對所有方法的相同圖像都使用了相同的掩碼。

(7)

式中:N為樣本的數(shù)量,W、H、C分別為圖像的寬度、高度以及通道數(shù),Iout為修復圖像,Igt為真實圖像。

(8)

(9)

式中:X1為修復圖像,X2為真實圖像,H、W分別為圖像的高度和寬度。

(10)

(11)

式中:μx、∑x分別為修復圖像在Inception Net-V3輸出的2 048維特征向量集合的均值和協(xié)方差矩陣,μg、∑g分別為真實圖像在Inception Net-V3輸出的2 048維特征向量集合的均值和協(xié)方差矩陣,tr為矩陣的跡。

定量分析結果如表1和表2所示。不難看出,本文模型在L1Loss、PSNR、SSIM和FID指標下都優(yōu)于CA模型、PConv模型、GatedConv模型、CRA模型。本文模型與MAT模型的定量分析結果如表3所示,不難看出,本文模型的L1Loss、PSNR、SSIM均優(yōu)于現(xiàn)階段人臉修復MAT模型,由此說明本文提出的模型與現(xiàn)有先進深度模型相比,具有更優(yōu)秀的人臉面部修復能力。

表1 CA模型、PConv模型和本文模型在CelebA-HQ測試集下的平均指數(shù)對比Tab.1 Average index comparison of CA model,PConv model and proposed model under CelebA-HQ test set

表2 GatedConv模型、CRA模型和本文模型在CelebA-HQ測試集下的平均指數(shù)對比Tab.2 Average index comparison of GatedConv model,CRA model and proposed model under CelebA-HQ test set

表3 MAT模型和本文模型在CelebA-HQ測試集下的平均指數(shù)對比Tab.3 Average index comparison between MAT model and proposed model under CelebA-HQ test set

5 結束語

本文提出了一種基于掩碼預測和多尺度上下文聚合的高分辨率人臉圖像修復模型。該模型由一個生成器和一個判別器組成。生成器由編碼器、多層MSCAM和解碼器組成。為了生成看似合理的內容,提出了MSCAM來構造生成器,融合了來自不同感受野的特征,且可以捕捉遙遠距離的上下文信息和感興趣的模式進行上下文推理。為了改善紋理合成,通過使用MP-PatchGAN判別器迫使它區(qū)分真實和生成的小塊(缺失區(qū)域)的紋理細節(jié)。此外,使用跳躍連接將編碼器中每一層卷積層的輸出與解碼器對應位置的輸入在通道維度上拼接,使得圖像的上下文信息向更高層分辨率特征圖傳播。實驗表明,本文模型在人臉圖像修復任務中能夠生成合理的內容和逼真的細粒度紋理。