基于雙重注意力機(jī)制的圖像修復(fù)方法*

王 威 李 眾

（江蘇科技大學(xué)電子信息學(xué)院 鎮(zhèn)江 212003）

1 引言

圖像修復(fù)是通過圖像中的已知信息來推測(cè)出缺失的像素值的過程，是計(jì)算機(jī)視覺領(lǐng)域和計(jì)算機(jī)圖形學(xué)的熱點(diǎn)研究方向之一。當(dāng)需要移除圖像中不需要的部分或是恢復(fù)被遮擋區(qū)域的時(shí)候，都需要進(jìn)行圖形修復(fù)。圖像修復(fù)的重點(diǎn)在于推測(cè)出的圖像在語(yǔ)義上和真實(shí)圖像要保持一致，并且在紋理細(xì)節(jié)上要逼真。

早期的圖像修復(fù)算法主要分為兩種，一種是基于圖像塊的算法［1~3］，通過在圖像中匹配合適的圖像塊并復(fù)制到缺失區(qū)域來實(shí)現(xiàn)圖像的修復(fù)；另一種是基于擴(kuò)散的算法［4~6］，將圖像內(nèi)容從缺失區(qū)域的邊界平穩(wěn)地傳播到缺失區(qū)域的內(nèi)部。這些傳統(tǒng)算法本質(zhì)上是利用圖像本身的冗余性，用圖像已知部分的信息來補(bǔ)全未知部分，但是無(wú)法獲得圖像高級(jí)別的語(yǔ)義。這些方法在應(yīng)對(duì)簡(jiǎn)單一致的背景修復(fù)任務(wù)時(shí)能取得不錯(cuò)的效果［7］，然而，當(dāng)缺失區(qū)域紋理復(fù)雜時(shí)，比如物體或是自然場(chǎng)景，它們都無(wú)法很好地完成修復(fù)任務(wù)。

得益于卷積神經(jīng)網(wǎng)絡(luò)（CNN）對(duì)于提取特征的高效［8］，以及生成對(duì)抗網(wǎng)絡(luò)［9］（GAN）取得的大量進(jìn)展，大量研究人員將其應(yīng)用在圖像修復(fù)領(lǐng)域。相比于傳統(tǒng)算法，基于CNN 和GAN 的方法通過大量的訓(xùn)練數(shù)據(jù)，可以學(xué)習(xí)到真實(shí)圖像中的特征以及抽象的語(yǔ)義特征，然后用于完成缺失部分的修復(fù)任務(wù)，并不像傳統(tǒng)算法是對(duì)圖像塊的復(fù)制，因此修復(fù)結(jié)果在上下文信息上更加連貫。Pathak 等［10］首次基于編碼器-解碼器機(jī)制提出Context Encoder 修復(fù)模型，其將圖像修復(fù)任務(wù)轉(zhuǎn)換為是有約束的圖形生成問題。通過編碼器進(jìn)行高級(jí)別的特征提取，用解碼器來預(yù)測(cè)缺失圖像的內(nèi)容，同時(shí)加入對(duì)抗網(wǎng)絡(luò)共同訓(xùn)練，以加強(qiáng)生成的像素與現(xiàn)有像素之間的一致性。實(shí)驗(yàn)表明該方法相比傳統(tǒng)方法在修復(fù)結(jié)果的語(yǔ)義一致性更好，能生成更加合理的圖像。但是其在紋理細(xì)節(jié)上還不夠精細(xì)，并且修復(fù)區(qū)域有很明顯的邊界。IIZUKA 等［11］在此基礎(chǔ)上將判別器模型修改為全局判別器和局部判別器相組合的方式來輔助生成器的訓(xùn)練，同時(shí)生成隨機(jī)二值掩膜，以進(jìn)行任意區(qū)域的修復(fù)。Chen 等［12］提出將生成器分為兩個(gè)階段，先將待修復(fù)圖片進(jìn)行粗略的修復(fù)，再將其輸出作為第二階段精細(xì)修復(fù)的輸入，獲得了更好的修復(fù)效果。這些方法［10~13］表明結(jié)合CNN 和GAN 的算法可以生成語(yǔ)義上合理的圖像，但是基于GAN的模型難以訓(xùn)練，并且直接對(duì)原始輸入使用CNN導(dǎo)致內(nèi)存占用量非常高，因此即使有較好的硬件支持，訓(xùn)練速度仍然很慢。隨著缺失區(qū)域的增加，修復(fù)質(zhì)量會(huì)迅速下降。

針對(duì)上述問題，為了改善模型對(duì)于圖像的上下文信息的理解能力，以及生成紋理更加清晰的圖像，本文設(shè)計(jì)了一種結(jié)合雙重注意力機(jī)制的兩階段修復(fù)方法：將生成器分成粗修復(fù)和精修復(fù)兩階段，同時(shí)在精修復(fù)階段的模型中加入雙重注意力機(jī)制，可以聚合圖像特征的空間相關(guān)性和通道相關(guān)性，使生成器更好地利用圖像的上下文信息。為了讓修復(fù)結(jié)果更加逼真，采用多尺度判別器，利用三個(gè)結(jié)構(gòu)相同，但輸入圖像尺度不同的判別器同時(shí)訓(xùn)練。在訓(xùn)練時(shí)，判別器的梯度會(huì)反向傳播給生成器，使生成器可以獲得不同尺度的梯度信息。

2 基于雙重注意力機(jī)制的圖像修復(fù)模型

2.1 雙重注意力機(jī)制

卷積運(yùn)算有一個(gè)顯著缺陷，其僅在局部相鄰區(qū)域工作，也由此會(huì)錯(cuò)失全局信息。不少研究只采用U-net 結(jié)構(gòu)去融合底層和高層的語(yǔ)義特征，但還是沒有綜合考慮各個(gè)位置的聯(lián)系和相關(guān)性。雙重注意力機(jī)制由空間注意力和通道注意力兩部分組成。

2.1.1 空間注意力

注意力機(jī)制有匹配和替換兩個(gè)階段。在匹配階段計(jì)算特征之間的注意力得分，在替換階段，通過注意力得分聚合完整區(qū)域的特征塊來替換缺失區(qū)域的特征塊。空間注意力得分在計(jì)算時(shí)首先將特征圖切分成3×3 大小的塊，然后利用余弦相似度來計(jì)算缺失部分的特征塊與完整區(qū)域的特征塊的相關(guān)性：

其中pi表示特征圖中完整區(qū)域的特征塊，pj表示特征圖中缺失部分的特征塊。再利用softmax 運(yùn)算得到每個(gè)特征塊的空間注意力得分：

其中N表示特征圖中完整區(qū)域的個(gè)數(shù)。最后在替換階段，通過空間注意力得分，聚合完整區(qū)域的特征塊來重建特征圖中的缺失區(qū)域：

其中α為重建系數(shù)，將重建的特征圖P?作為新的特征圖傳遞給通道注意力機(jī)制，進(jìn)行第二次的重建，以更好地利用全局的語(yǔ)義信息。

2.1.2 通道注意力

特征圖中的每一層都可以看作是卷積核對(duì)一個(gè)特定特征的響應(yīng)，而不同的語(yǔ)義響應(yīng)相互關(guān)聯(lián)。通過計(jì)算通道之間的相互依賴性，可以強(qiáng)調(diào)相互依賴的特征圖，改善特定語(yǔ)義的特征表示。本文的通道注意力采用自注意力模型，將各層特征圖轉(zhuǎn)置之后進(jìn)行矩陣乘法來計(jì)算各層之間的相關(guān)性：

其中Ai表示第i層特征圖，Aj表示第j層特征圖。利用softmax 運(yùn)算來計(jì)算各層通道注意力分?jǐn)?shù)：

其中C表示特征圖的通道數(shù)。最后通過融合各層通道注意力分?jǐn)?shù)來修正特征圖：

其中β為重建系數(shù)。空間注意力機(jī)制的輸出作為通道注意力機(jī)制的輸入，相較于只使用單個(gè)注意力機(jī)制，雙重注意力機(jī)制可以獲取更豐富的語(yǔ)義信息。

2.2 兩階段生成器模型

生成模型的主要工作原理是模型根據(jù)輸入的待修復(fù)圖像，輸出完整的在語(yǔ)義上一致的圖像。然后通過裁剪將生成的圖像粘貼到待修復(fù)區(qū)域。生成器的網(wǎng)絡(luò)架構(gòu)如圖1 所示。本文采用了兩階段網(wǎng)絡(luò)架構(gòu)，其中第一階段粗略的生成缺失內(nèi)容，而第二階段則生成出更精細(xì)的結(jié)果。

圖1 圖像修復(fù)流程圖

對(duì)于第一階段，首先生成隨機(jī)二值掩膜，同時(shí)處理圖像來模擬任意區(qū)域的缺失部分，然后將缺失圖像和對(duì)應(yīng)的二值掩膜作為第一階段的輸入進(jìn)行訓(xùn)練，完成粗略的修復(fù)。第一階段模型參考U-Net［15］全卷積結(jié)構(gòu)，去掉其中的池化層，采用步幅卷積來實(shí)現(xiàn)降采樣，避免池化操作對(duì)于信息的丟失，同時(shí)將其中的有效卷積修改成同維卷積。將第一階段的輸出以及對(duì)應(yīng)的二值掩膜作為第二階段的輸入。第二階段中通過雙重注意力機(jī)制來對(duì)高層特征圖處理，改善特定語(yǔ)義的特征表示。在圖像修復(fù)任務(wù)中，模型感受野的尺寸應(yīng)該盡可能的大［16］。因此本文采用了擴(kuò)充卷積用來擴(kuò)大模型的感受野，同時(shí)加入殘差塊結(jié)構(gòu)和跳躍連接，在一定程度上也可以緩解生成器由于模型過深而導(dǎo)致的模型退化和梯度消失［17］。擴(kuò)張卷積可以在不改變特征圖尺寸的情況下增大感受野，但同時(shí)導(dǎo)致不能讓所有數(shù)據(jù)參與計(jì)算，體現(xiàn)在特征圖上就是卷積中心點(diǎn)的不連續(xù)［18］，因此在設(shè)計(jì)擴(kuò)張卷積結(jié)構(gòu)時(shí)將擴(kuò)張率改成鋸齒狀。

2.3 多尺度判別器模型

在圖像修復(fù)任務(wù)中，最重要的是對(duì)缺失部分像素值的預(yù)測(cè)。在對(duì)抗學(xué)習(xí)中，判別器通過判別圖像是否與原圖相似來約束生成器的訓(xùn)練。好的判別器可以有效地保持修復(fù)圖像整體的語(yǔ)義一致性，增強(qiáng)修復(fù)區(qū)域的紋理細(xì)節(jié)。

本文使用多尺度判別器來輔助生成器訓(xùn)練。多尺度判別器由三個(gè)相同結(jié)構(gòu)的判別器構(gòu)成，通過最大池化處理圖像來實(shí)現(xiàn)降采樣，以獲得不同尺度的圖像分別作用于三個(gè)判別器。小尺度的判別器模型能獲得的更大的感受野，使得生成圖像的語(yǔ)義一致性就好。大尺度的判別器模型能學(xué)習(xí)到更精細(xì)的紋理結(jié)構(gòu)，使得生成圖像更加清晰。

為了能更好地約束GAN 的訓(xùn)練，對(duì)判別器每一層的卷積核參數(shù)矩陣W施加譜范數(shù)歸一化（Spectral Normalization，SN），參數(shù)矩陣更新公式計(jì)算方法如式（7）所示：

其中σmax(W)是WTW的最大特征值。譜范數(shù)歸一化嚴(yán)格實(shí)現(xiàn)利普希茨連續(xù)性約束，使得判別器對(duì)于輸入擾動(dòng)具有更好的穩(wěn)定性，防止梯度異常，從而使訓(xùn)練過程更容易收斂。Miyato 等［14］證明，使用譜范數(shù)歸一化的判別器模型比使用梯度裁剪的要更穩(wěn)定。

2.4 損失函數(shù)

本文修復(fù)模型的損失函數(shù)包括對(duì)抗損失和重建損失兩個(gè)部分。對(duì)抗損失通過對(duì)抗訓(xùn)練使生成器生成更加合理的圖像，重建損失可以強(qiáng)化生成圖像的紋理細(xì)節(jié)與原始圖像之間的一致性。

本文的對(duì)抗損失采用的是WGAN 損失［19］，因?yàn)橐呀?jīng)使用了譜范數(shù)歸一化來約束模型波動(dòng)，因此本文中不使用梯度懲罰。WGAN 損失采用的是Wasserstein 距離用來衡量生成數(shù)據(jù)和真實(shí)數(shù)據(jù)之間分布的差異，在確保了生成樣本的多樣性的同時(shí)，改善了原始GAN 訓(xùn)練不穩(wěn)定的問題。依據(jù)WGAN 的目標(biāo)函數(shù)設(shè)計(jì)多尺度判別器的損失函數(shù)Ld和對(duì)抗損失Ladv：

其中Pg和Pr分別表示通過生成器生成的數(shù)據(jù)的概率分布和真實(shí)數(shù)據(jù)的概率分布。通過最小化生成器損失以及最大化判別器損失來優(yōu)化網(wǎng)絡(luò)。

WGAN 在訓(xùn)練生成器G 的同時(shí)聯(lián)合訓(xùn)練判別器D為生成器提供梯度，通過兩個(gè)網(wǎng)絡(luò)之間相互對(duì)抗，使G盡可能地逼近真實(shí)數(shù)據(jù)分布。不同于原始GAN 的判別器是做二分類任務(wù)，WGAN 的判別器屬于回歸任務(wù)，因此最后一層需要去掉sigmoid 運(yùn)算。本文采用L1 范數(shù)損失作為重建損失，相較于L2范數(shù)損失，能恢復(fù)更多的邊緣信息，避免紋理過于平滑。

本文的修復(fù)模型采用兩階段訓(xùn)練，在訓(xùn)練第一階段網(wǎng)絡(luò)時(shí)，只使用重建損失來獲得粗略的修復(fù)結(jié)果。第二階段強(qiáng)化訓(xùn)練，使用重建損失和對(duì)抗損失相結(jié)合的方式進(jìn)行訓(xùn)練。

3 實(shí)驗(yàn)與分析

3.1 訓(xùn)練

本文使用Place 2 數(shù)據(jù)集［20］作為訓(xùn)練數(shù)據(jù)，其包含許多場(chǎng)景類別，例如臥室、街道、猶太教堂、峽谷等。數(shù)據(jù)集由1000 萬(wàn)張圖像組成，其中每個(gè)場(chǎng)景類別包含5000 張訓(xùn)練圖像和500 張測(cè)試圖像。每張?jiān)紙D像的尺寸為256×256，掩膜為128×128的矩形。

本文在基于tensorflow-gpu 1.14.0 框架的異構(gòu)平臺(tái)上進(jìn)行試驗(yàn)，顯卡為NVIDIA GeForce GTX1080，對(duì)應(yīng)的cuda 版本為10.0，batch size 大小設(shè)置為8，對(duì)抗訓(xùn)練輪數(shù)為1000000 次。為了消除奇異樣本數(shù)據(jù)導(dǎo)致的不良影響，將圖像的像素值歸一化到［-1，1］之間作為生成器模型的輸入。

3.2 圖像修復(fù)結(jié)果評(píng)估

圖2 展示了本文的方法與PatchMatch（PM）方法［6］和Context Encoder（CE）方法［9］修復(fù)結(jié)果的比較。圖2（c）為PM 方法修復(fù)結(jié)果，可以看到其修復(fù)結(jié)果紋理十分清晰，但是這種方法缺少對(duì)圖像的語(yǔ)義理解，第一張圖和第三張圖出現(xiàn)了很不合理的部分。圖2（d）為CE 方法修復(fù)結(jié)果，其對(duì)于全局的語(yǔ)義有一定的理解，相較于PM方法，修復(fù)結(jié)果在結(jié)構(gòu)一致性上會(huì)更好，但是修復(fù)區(qū)域十分模糊，并且有很明顯的修復(fù)邊界，無(wú)法達(dá)到圖像修復(fù)任務(wù)的要求。圖2（e）為本文方法修復(fù)結(jié)果，可以看到本文的方法可以更好地利用周圍的紋理和結(jié)構(gòu)，能恢復(fù)出更多的有效信息。同時(shí)生成的圖像也在視覺上更真實(shí)，不會(huì)產(chǎn)生明顯的修復(fù)邊界和偽影。

圖2 與其他修復(fù)方法對(duì)比

為了量化圖像修復(fù)結(jié)果，本文給定評(píng)價(jià)指標(biāo)為峰值信噪比［21］（PSNR）、結(jié)構(gòu)相似性（SSIM）和平均誤差（L1 loss）。PSNR［22］通過像素級(jí)別的誤差來評(píng)估生成圖像與原圖像的差異，其計(jì)算方法如式（11）。

其中，MSE 是原圖像與修復(fù)圖像之間均方誤差。PSRN 值越大，圖像之間的像素差異越小。SSIM 分別從亮度、對(duì)比度和結(jié)構(gòu)三個(gè)方面來評(píng)估生成圖像與原圖像的相似性，計(jì)算方法如式（12）：

其中μ是圖像的平均值，σ是圖像的方差。C1和C2是用來維持穩(wěn)定的常數(shù)。SSIM 值越大整體結(jié)構(gòu)約相似。L1 loss是輸入和輸出之間絕對(duì)差值之和，值越小，表示圖像失真越小。計(jì)算方法如式（13）：

表1 統(tǒng)計(jì)了本文方法的修復(fù)結(jié)果與各個(gè)方法修復(fù)結(jié)果幾組數(shù)據(jù)的平均值，可以看到本文方法在三個(gè)評(píng)價(jià)指標(biāo)上都優(yōu)于其他方法。

表1 圖像修復(fù)評(píng)價(jià)指標(biāo)

4 結(jié)語(yǔ)

本文提出了一種基于生成對(duì)抗網(wǎng)絡(luò)的圖像修復(fù)方法，利用雙重注意力機(jī)制和殘差網(wǎng)絡(luò)來增加模型對(duì)于圖像的上下文信息的理解與利用。同時(shí)設(shè)計(jì)了一種基于多尺度判別器的對(duì)抗訓(xùn)練模型，給判別器模型提供不同尺度的梯度信息，幫助模型生成更加清晰的圖像。實(shí)驗(yàn)結(jié)果表明，相比于現(xiàn)階段的方法，本文的方法充分利用圖像中的上下文信息，能生成視覺上更加合理的圖像，同時(shí)與周圍區(qū)域的紋理一致性上會(huì)更好。通過PSNR、SSIM 和l1 損失對(duì)修復(fù)結(jié)果進(jìn)行評(píng)估也表明本文的方法有更高的可信度。