融合邊緣檢測和自注意力的圖像修復方法①

李維乾,張曉文

(西安工程大學 計算機科學學院,西安 710600)

圖像修復是計算機視覺領域的一個重要任務.它在數字藝術品修復、公安刑偵面部修復等種種實際場景中被廣泛應用.圖像修復的核心挑戰在于為缺失區域合成視覺逼真和語義合理的像素,要求合成的像素與原像素具有一致性.

傳統的圖像修復技術有基于結構和紋理兩種方法.基于結構的圖像修復算法具有代表性的是Bertalmio 等[1]提出的BSCB 模型和Shen 等[2]提出的基于曲率擴散的修復模型CDD.基于紋理的修復算法中具有代表性的有Criminisi 等[3]提出的基于patch的紋理合成算法.這兩種傳統的修復算法可以修復小塊區域的破損,但是在破損區域越來越大時,修復效果則直線下降,并且修復結果存在圖像模糊、結構扭曲、紋理不清晰和視覺不連貫等問題.

近年來隨著硬件設備等計算能力的不斷提升,以及深度學習技術在圖像翻譯[4,5]、圖像超分辨率[6]、圖像修復[7-9]等計算機視覺領域的迅速發展,采用深度學習技術的修復方法能夠捕獲圖像的高層語義信息,與傳統的修復方法相比,具有良好的修復效果.現有基于深度學習的修復方法包括卷積神經網絡的方法和生成式對抗網絡(Generative Adversarial Networks,GAN)的方法,其中基于GAN的方法在圖像修復領域的應用比較出色.該技術采用對抗訓練的方式來幫助模型生成需要的樣本圖像,采用GAN 方法能夠提高數據質量,提升學習效率,但存在修復區域與周圍區域語義不相關、邊界模糊、訓練不穩定等問題,因此相關學者對基于GAN 網絡的圖像修復方法進行了改進.如Pathak 等[10]提出的Context Encoder 模型,該模型使用對抗損失訓練一個上下文編碼器預測圖像的受損區域,但使用該方法修復的圖像比較模糊且不一致;Iizuka 等[11]提出了雙鑒別器的生成對抗模型,該模型通過增設局部鑒別器跟蹤受損區域保證了圖像的一致性;Liu 等[12]提出的連貫性語義注意力方法,通過搜索已知區域中最相似的特征patch 填滿缺失區域,增加patch 判別器進行一致性優化,保證了全局一致性;為了得到符合語義的圖像,Yu 等[13]提出了基于上下文注意力的兩階段模型,加入上下文注意力模塊,實現提取局部區域的有效特征,但是它不能處理任意形狀的受損區域.還有一些利用其它技術的研究,如Nazeri 等[14]提出了基于邊緣連接的修復模型,它應用隨機生成二值掩碼解決了不能處理任意形狀受損區域的問題,通過邊緣連接網絡生成圖像邊緣信息,有效解決了圖像邊界模糊的問題;Ren 等[15]提出的結構感知修復模型,通過邊緣圖像來表示圖像的全局結構,雖然基于邊緣特征的方法有效地描繪了輪廓信息,但是它有時無法準確地描繪高度紋理化區域.

盡管上述方法在圖像修復任務中表現了出色的性能,但還存在以下不足:首先圖像的高頻紋理信息無法準確修復;其次雖然保證了圖像在視覺上的一致性和連續性,但是有些區域的邊界和紋理存在修復模糊的現象.

因此,針對修復后圖像邊界模糊、紋理不清晰、視覺效果差的問題,本文提出了融合邊緣檢測和自注意力機制的生成對抗修復模型.該模型由邊緣補全網絡和紋理修復網絡構成,首先利用邊緣補全網絡填補圖像的邊緣,得到邊緣補全圖像;其次通過紋理修復網絡聯合邊緣補全圖像修復缺失的紋理信息.由于卷積神經網絡主要處理局部區域的特征,僅使用卷積神經網絡很難建立長相關性模型,而自注意力機制可以處理圖像中全局范圍的特征,可以建立長范圍、多層次的依賴關系.因此本文將卷積神經網絡與自注意力機制相結合設計出紋理修復網絡來生成缺失區域的紋理信息,最終使模型生成的圖像與原圖像在邊緣、紋理結構、視覺上高度逼真.

1 生成式對抗網絡

生成式對抗網絡由Goodfellow 于2014年提出[16].GAN 由一個生成器 Generator和一個鑒別器 Discriminator 組成,生成器G的輸入是一個隨機噪聲z,通過這個噪聲來生成“以假亂真”的圖像G(z).鑒別器D的輸入是真實圖像和生成圖像的混合,它的目的是判別輸入圖像的真實性.網絡通過對抗訓練讓生成器和鑒別器同時得到增強,訓練時,生成器G用來生成逼真的“假”圖像去混淆D的判斷,鑒別器D要判別出生成器G生成的圖像為假,同時給生成器一個反饋來更新網絡的參數,使生成器生成質量更高的圖像,鑒別器的判別能力也會隨之提高.GAN的基本網絡結構如圖1所示.

圖1 GAN的基本網絡結構

生成對抗網絡的目標是最小化生成器生成的圖像與真實圖像概率分布之間的差距.因此生成式對抗網絡的目標函數可以表示為:

式中,G和D分別表示生成器與鑒別器的可微函數,x是真實數據樣本,z是隨機噪聲矢量,G(z)是生成器的生成數據.E代表數學期望值.在同時訓練生成器和鑒別器的情況下,GAN 訓練的目標是令判別器D的輸出D(G(z))趨近于0,D(x)趨近1,然后將結果反饋給生成網絡,生成器調整自身參數,使生成器下一次生成更加逼真的圖像.訓練中一次次迭代使生成器和鑒別器最終達到一個“平衡點”,即生成器生成的圖片已經“以假亂真”,令鑒別器無法判別圖像的真偽.

2 融合邊緣檢測和自注意力的修復模型

為了解決圖像修復技術中存在的邊界模糊、紋理不清晰、視覺效果差的問題,本文提出了融合邊緣檢測和自注意力機制的生成對抗修復模型.如圖2所示,整個模型包含邊緣補全網絡和紋理修復網絡兩個部分.首先檢測真實圖像和受損圖像的邊緣,得到其邊緣圖像,以真實圖像的邊緣為標簽,受損圖像及其邊緣圖像為輸入訓練邊緣補全網絡,補全缺失區域的輪廓信息,得到邊緣補全圖像;其次,以邊緣補全圖像和受損圖像作為輸入,對融合注意力機制的紋理修復網絡進行訓練,使生成器G2生成缺失區域的精確紋理信息,從而使紋理修復網絡完成缺失區域紋理的精確修復;最后將邊緣補全網絡的生成器G1和紋理修復網絡的生成器G2進行級聯,以受損圖像及其邊緣圖像作為模型的輸入,以真實圖像為標簽進行聯合訓練,實現端到端的圖像修復.

圖2 融合邊緣檢測和自注意力機制的生成對抗修復模型

2.1 邊緣補全

為了在圖像修復中提供清晰的邊界和結構,本文只提取物體的輪廓信息.傳統的基于Sobel和Canny的檢測方法只考慮顏色、光線和梯度局部變化.在復雜的場景中,它們是無法滿足需要的.因此本文選擇整體嵌套邊緣檢測(Holistically-nested Edge Detection,HED)[17]方法.它在處理的是整個圖像,并能夠提取出圖像的高層邊界信息.本文使用已經訓練好的HED 模型提取圖像的邊緣.

圖像邊緣的生成分為兩個步驟,第一是從受損圖像中提取邊緣信息,第二是補全缺失的邊緣.如圖2左半部分所示,邊緣補全網絡包含一個生成器G1和一個鑒別器D1.生成器G1由對圖像進行兩次下采樣的編碼器,隨后的8 個殘差塊和將圖像上采樣回原始大小的解碼器組成.邊緣補全網絡的輸入由受損圖像及其邊緣圖像組成,通過卷積和標準化處理,提取出圖像深層次紋理邊緣信息,通過對邊緣特征的學習訓練,最終讓生成器G1預測出缺失區域的邊緣.為了提取深層次的邊緣紋理特征,在卷積和反卷積之間增加了殘差塊,避免了由于網絡深度過深而導致的梯度彌散問題.

鑒別器D1是判別邊緣修復是否準確的網絡.本文利用HED的邊緣檢測方法提取出真實圖像的邊緣信息,將其作為網絡訓練的正樣本,然后,將生成器G1生成的樣本作為負樣本和正樣本進行混合輸入鑒別器D1,判別其真實性,通過不斷優化損失函數來提升鑒別器D1的判別能力,生成器G1的生成能力隨著鑒別器D1判別能力的提升而不斷上升,促使生成器G1生成的邊緣圖像與原圖的邊緣圖像趨于一致.

Igt表示真實圖像,其邊緣圖用Bgt表示.在邊緣補全網絡中,使用受損圖像Iin及其邊緣圖像作為輸入,其邊緣圖像為=Bgt⊙(1-M),其中M為受損圖像Iin的掩碼,該變量為一個二值矩陣(缺失區域為1,背景為0);在此,⊙表示矩陣乘法.將受損圖像Iin和它的邊緣圖像作為生成器G1的輸入,那么邊緣補全網絡生成的缺失區域的邊緣圖像是:

其中,Bpred為邊緣補全圖像.鑒別器D1以Bgt和Bpred作為輸入來判斷邊緣補全圖像是否真實.用對抗損失作為目標函數對網絡進行訓練,通過優化目標函數不斷提升鑒別器和生成器的能力來實現圖像邊緣信息的補全.

2.2 紋理修復

如圖2右半部分所示,紋理修復網絡由生成器G2和雙鑒別器D2、D3組成,生成器G2由對圖像進行兩次下采樣的編碼器,引入自注意力機制的殘差塊和將圖像上采樣回原始大小的解碼器組成,其輸入是邊緣補全圖像和受損圖像的結合.該網絡將自注意力機制與卷積神經網絡進行結合來提取輸入圖像深層次特征.自注意力機制會利用所有位置的特征來生成圖像某一位置的細節.通過自注意力機制計算得到輸入圖像的注意力特征圖,然后將注意力特征圖與提取的卷積特征圖進行連接輸入到網絡中繼續進行訓練,通過不斷優化模型提升生成器的能力,最后促使生成器生成與原圖像趨于一致的圖像.

鑒別器由全局鑒別器D2和局部鑒別器D3構成雙鑒別器機制.全局鑒別器D2的作用是判別修復后圖像的整體真實性,它的作用集中在圖像的全部區域,它監督生成器生成視覺一致的圖像;局部鑒別器D3用于判別圖像受損區域的真實性.鑒別器的輸入為生成器的生成圖像和真實圖像的混合,通過優化對抗損失和重構損失函數提升鑒別器的判別能力,讓鑒別器與生成器進行抗衡,來迫使生成器生成與原圖紋理細節一致的圖像.

紋理修復網絡是一個加入自注意力機制的生成式對抗網絡.生成器G2用于生成生動的圖像紋理信息.鑒別器D2和D3分別用于判別生成的圖像全局和局部區域是否為真.將邊緣補全圖像Bpred和受損圖像Iin一起作為紋理修復網絡的輸入,最終生成的完整圖像是:

其中,Iout為紋理修復網絡的輸出結果,它是缺失區域的完整修復結果.

2.3 自注意力機制

在圖像修復任務中,最終目標是訓練一個生成器網絡,使其生成與真實圖像高層紋理細節盡可能相似的圖像.為了實現這一點,引入自注意力機制,它可以更好地利用全局特征信息對圖像進行重建.

卷積操作主要處理局部區域間的信息,因此多個卷積層在建立全局依賴性上計算效率非常低.自注意力機制提供了全局的感受野,它會學習某個像素點和其它所有位置像素點的關系,使得每個位置的生成不再是僅僅依賴這個位置臨近的其他位置,而是引入了距離更遠且具有更大信息量的位置,每個位置的特征都由其它所有位置的加權和來計算得到.因此,本文借鑒文獻[18]思路,在紋理生成網絡的殘差塊中引入自注意力機制,構建自注意力機制的生成式對抗網絡,通過該網絡獲取圖像的全局特征信息,從而更好地重建出圖像的紋理細節,如圖3所示.

圖3 自注意力機制

圖3中,從殘差塊第二層中提取的卷積特征圖x分別經過兩個1 ×1的卷積層變換為兩個特征空間f(x),g(x)去計算注意力圖.f(x)用 來提取像素特征,g(x)用來提取全局特征信息,其中f(x)=Wf x,g(x)=Wgx.

其中,sij=f(xi)Tg(xj).βj,i表示模型在合成第j個區域時對第i個位置的關注程度.N表示前一層特征圖的特征位置數.注意力特征圖的輸出是o=(o1,o2,···,oj,···,oN).

其中,h(xi)=Whxi,v(xi)=Wvxi.此外再將自注意力機制的輸出乘以一個比例參數,并添加到輸入自注意力機制的卷積特征圖中,因此,最終輸出為:

其中,yi是自注意力機制的最終輸出結果,將其作為下一個殘差塊的輸入.γ是一個可學習的標量,它的初始值為0,引入γ 讓網絡首先依賴局部相鄰的線索,然后逐漸學習給非局部線索分配更多的權重,這樣做的原因是先讓模型學習簡單的任務,然后逐漸增加任務的復雜性.

2.4 損失函數

傳統的生成式對抗網絡在訓練時使用對抗損失作為目標函數,但存在訓練不穩定的問題.本文使用上下文注意力模型[13]中基于WGAN-GP的修改版本作為訓練網絡的目標函數,促使網絡生成視覺逼真的圖像.WGAN-GP 損失在圖像生成任務方面的性能比現有的GAN 損失性能強.

WGAN 使用EM 距離W(Pr,Pg)來比較真實數據和生成數據的分布.WGAN的目標函數是:

其中,d是1-Lipschitz 函數的集合,Pr是真實數據的分布,Pg是=G(z)的模型分布,z是生成器的輸入.對于圖像修復來說,只需要修復圖像缺失的區域,所以梯度懲罰項為:

(1)邊緣補全的損失.本文將WGAN-GP 損失應用到模型的兩個部分中,對邊緣補全網絡來說,它的對抗損失定義為:

其中,PBpred是邊緣補全圖像的數據分布,PBgt是真實邊緣圖像數據的分布,是從生成數據分布和真實數據分布采樣的點之間的直線中采樣的.

(2)紋理修復的損失.紋理修復網絡的輸出為修復結果Iout.它的對抗損失定義為:

其中,PIout是生成的修復圖像的分布,PIgt是真實圖像的分布.

WGAN-GP 損失和l1重構損失聯合使用時更容易訓練,優化過程更穩定,因此在紋理修復網絡中使用重構損失和WGAN-GP 損失組合作為目標函數對網絡進行訓練.紋理修復網絡的重構損失被定義為:

其中,N是樣本的數量,W,H,C分別是圖像的寬度、高度和通道數.那么紋理修復網絡訓練的目標損失函數為:

其中,λ1和λ2分別是重構損失和對抗損失的參數,在實驗中,λ1=1,λ2=4.

3 實驗結果及分析

3.1 數據集

為了驗證本文圖像修復方法的效果,本文實驗中用Place2和CelebA 數據集來訓練和評估本文的兩階段模型,CelebA 人臉數據集包含10 177 個名人身份的202 599 張人臉圖像.Place2是一個場景圖片數據集,它包含1 千萬張圖片,400 多個不同種類的場景.將掩碼設置為不規則覆蓋,覆蓋范圍設置整幅圖像的10-20%.

3.2 參數設置

本文的實驗平臺為Windows10 系統、Python3.6和PyTorch1.0.0 結合的編程環境.CPU 信息為Intel(R)Core(TM)i7-8550U CPU,主頻2.0 GHz,內存為16 GB,NVIDIA Titan XP GPU.在訓練中使用 256×256的圖像,batch size 設置為32.使用Adam 優化器[19]來優化模型,它的參數設置為β1=0.9和β2=0.999.本文設置學習率為10-3來訓練G1和G2生成器直到收斂為止,鑒別器以生成器學習率的0.1 倍進行訓練.

3.3 實驗結果及分析

為了分析本文方法在多個場景的應用效果,在Place2和CelebA 數據集上進行效果評估.并與近年來基于生成式對抗網絡具有代表性的圖像修復方法進行比較(上下文注意力的生成式對抗模型[13]和邊緣連接的生成式對抗模型[14]).本文用峰值信噪比(PSNR)和結構相似性指數(SSIM)對實驗結果進行定量評估,以此來表示修復后圖像和真實圖像的差距.

峰值信噪比PSNR是一種評價圖像的客觀標準,它直接衡量像素值的差異,數值越大表示失真越小.假設兩張圖像分別是X1和X2,計算公式如下:

結構相似性指數(SSIM)是衡量圖像相似度的指標.它是一個0-1 之間的數,這個數越大,表示修復圖像和真實圖像之間越相似,圖像更逼真.假設兩張圖像分別是X1和X2,計算公式如下:

其中,μX1和μX2分別為X1和X2的均值,σ2X1和σ2X2分別為X1和X2的方差,σX1X2為X1和X2的協方差,c1和c2為兩個常數,避免除零.

本文的模型在CelebA和Place2 兩個數據集上進行訓練和評估,通過大量數據對模型進行訓練,最后進行測試并計算該模型在不同數據集上的PSNR和SSIM平均值,并與基于上下文注意力的生成式對抗模型和基于邊緣連接的生成式對抗模型進行比較,以上3 種修復方法在數據集上的PSNR和SSIM平均數值對比如表1、表2所示.

表1 不同修復方法在CelebA 上的PSNR和SSIM平均指數對比

表2 不同修復方法在Place2 上的PSNR和SSIM平均指數對比

由表1和表2可以看出本文方法的PSNR和SSIM平均值比其他兩種方法的平均值高,實驗證明本文的算法在CelebA和Place2 數據集上的修復結果要優于其余兩種修復方法.

本文的方法在Place2和CelebA 數據集上的修復結果如圖4所示,由圖可以看出修復后得到的圖像和真實圖像在顏色、結構、風格上基本一致,且圖像的紋理細節構建比較準確和清晰;缺失區域與背景區域之間邊界自然清晰,視覺效果一致,沒有修復過的痕跡.這表明本文的方法對于圖像修復任務有明顯的效果.

不同修復算法在Place2和CelebA 數據集上的修復結果如圖5所示,由圖5可以看出,在CelebA 數據集的圖像中,對比人面部表情以及五官的修復效果來看,本文的方法修復效果最好;可以看到圖5(c)(上下文注意力的方法)中,缺失區域的修復在視覺上比較模糊;在圖5(d)(邊緣連接的方法)中,可以明顯看到修復的結果出現了偏差,從圖中可以看出修復后的人臉面部稍微扭曲,并且出現五官與原圖偏差的現象.

圖4 本文的修復方法結果圖

在Place2 數據集的圖像中,從房屋窗戶和圓錐建筑來看,本文的方法修復效果比較好,窗戶和圓錐建筑的基本輪廓還有紋理與原圖都是最接近的;圖5(c)中,可以看到屋頂和煙囪都比較模糊,圓錐建筑修復結果比較好;在圖5(d)中,可以看到窗戶和圓錐建筑基本輪廓不清晰,修復結果比較模糊.最后一幅圖像中,上下文注意力和邊緣連接的修復視覺上比較模糊,且河水出現了扭曲,說明修復出現了偏差.總體來說,本文方法的修復結果最佳,在一定程度上緩解了圖像邊界模糊的現象,視覺效果上比較一致、連貫,缺失區域的紋理修復準確且清晰.

4 結束語

本文提出了一種由邊緣到精細的兩階段模型,引入邊緣檢測和自注意力機制來解決圖像修復任務中存在的視覺效果差、邊界模糊和紋理不清晰等問題.本文基于生成式對抗網絡的模型,引入邊緣檢測得到受損圖像的邊緣補全圖,在紋理修復網絡中加入自注意力機制,利用自注意力具有全局視野的特點,讓紋理修復網絡生成了更加精確的紋理圖像.通過實驗結果表明,本文的方法在不同場景圖像的修復任務中,能夠生成邊界清晰,紋理精確,視覺連貫的圖像,表現了出色的性能.