基于雙重注意力機制的圖像修復方法*

王 威 李 眾

（江蘇科技大學電子信息學院 鎮江 212003）

1 引言

圖像修復是通過圖像中的已知信息來推測出缺失的像素值的過程，是計算機視覺領域和計算機圖形學的熱點研究方向之一。當需要移除圖像中不需要的部分或是恢復被遮擋區域的時候，都需要進行圖形修復。圖像修復的重點在于推測出的圖像在語義上和真實圖像要保持一致，并且在紋理細節上要逼真。

早期的圖像修復算法主要分為兩種，一種是基于圖像塊的算法［1~3］，通過在圖像中匹配合適的圖像塊并復制到缺失區域來實現圖像的修復；另一種是基于擴散的算法［4~6］，將圖像內容從缺失區域的邊界平穩地傳播到缺失區域的內部。這些傳統算法本質上是利用圖像本身的冗余性，用圖像已知部分的信息來補全未知部分，但是無法獲得圖像高級別的語義。這些方法在應對簡單一致的背景修復任務時能取得不錯的效果［7］，然而，當缺失區域紋理復雜時，比如物體或是自然場景，它們都無法很好地完成修復任務。

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

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

2 基于雙重注意力機制的圖像修復模型

2.1 雙重注意力機制

卷積運算有一個顯著缺陷，其僅在局部相鄰區域工作，也由此會錯失全局信息。不少研究只采用U-net 結構去融合底層和高層的語義特征，但還是沒有綜合考慮各個位置的聯系和相關性。雙重注意力機制由空間注意力和通道注意力兩部分組成。

2.1.1 空間注意力

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

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

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

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

2.1.2 通道注意力

特征圖中的每一層都可以看作是卷積核對一個特定特征的響應，而不同的語義響應相互關聯。通過計算通道之間的相互依賴性，可以強調相互依賴的特征圖，改善特定語義的特征表示。本文的通道注意力采用自注意力模型，將各層特征圖轉置之后進行矩陣乘法來計算各層之間的相關性：

其中Ai表示第i層特征圖，Aj表示第j層特征圖。利用softmax 運算來計算各層通道注意力分數：

其中C表示特征圖的通道數。最后通過融合各層通道注意力分數來修正特征圖：

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

2.2 兩階段生成器模型

生成模型的主要工作原理是模型根據輸入的待修復圖像，輸出完整的在語義上一致的圖像。然后通過裁剪將生成的圖像粘貼到待修復區域。生成器的網絡架構如圖1 所示。本文采用了兩階段網絡架構，其中第一階段粗略的生成缺失內容，而第二階段則生成出更精細的結果。

圖1 圖像修復流程圖

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

2.3 多尺度判別器模型

在圖像修復任務中，最重要的是對缺失部分像素值的預測。在對抗學習中，判別器通過判別圖像是否與原圖相似來約束生成器的訓練。好的判別器可以有效地保持修復圖像整體的語義一致性，增強修復區域的紋理細節。

本文使用多尺度判別器來輔助生成器訓練。多尺度判別器由三個相同結構的判別器構成，通過最大池化處理圖像來實現降采樣，以獲得不同尺度的圖像分別作用于三個判別器。小尺度的判別器模型能獲得的更大的感受野，使得生成圖像的語義一致性就好。大尺度的判別器模型能學習到更精細的紋理結構，使得生成圖像更加清晰。

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

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

2.4 損失函數

本文修復模型的損失函數包括對抗損失和重建損失兩個部分。對抗損失通過對抗訓練使生成器生成更加合理的圖像，重建損失可以強化生成圖像的紋理細節與原始圖像之間的一致性。

本文的對抗損失采用的是WGAN 損失［19］，因為已經使用了譜范數歸一化來約束模型波動，因此本文中不使用梯度懲罰。WGAN 損失采用的是Wasserstein 距離用來衡量生成數據和真實數據之間分布的差異，在確保了生成樣本的多樣性的同時，改善了原始GAN 訓練不穩定的問題。依據WGAN 的目標函數設計多尺度判別器的損失函數Ld和對抗損失Ladv：

其中Pg和Pr分別表示通過生成器生成的數據的概率分布和真實數據的概率分布。通過最小化生成器損失以及最大化判別器損失來優化網絡。

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

本文的修復模型采用兩階段訓練，在訓練第一階段網絡時，只使用重建損失來獲得粗略的修復結果。第二階段強化訓練，使用重建損失和對抗損失相結合的方式進行訓練。

3 實驗與分析

3.1 訓練

本文使用Place 2 數據集［20］作為訓練數據，其包含許多場景類別，例如臥室、街道、猶太教堂、峽谷等。數據集由1000 萬張圖像組成，其中每個場景類別包含5000 張訓練圖像和500 張測試圖像。每張原始圖像的尺寸為256×256，掩膜為128×128的矩形。

本文在基于tensorflow-gpu 1.14.0 框架的異構平臺上進行試驗，顯卡為NVIDIA GeForce GTX1080，對應的cuda 版本為10.0，batch size 大小設置為8，對抗訓練輪數為1000000 次。為了消除奇異樣本數據導致的不良影響，將圖像的像素值歸一化到［-1，1］之間作為生成器模型的輸入。

3.2 圖像修復結果評估

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

圖2 與其他修復方法對比

為了量化圖像修復結果，本文給定評價指標為峰值信噪比［21］（PSNR）、結構相似性（SSIM）和平均誤差（L1 loss）。PSNR［22］通過像素級別的誤差來評估生成圖像與原圖像的差異，其計算方法如式（11）。

其中，MSE 是原圖像與修復圖像之間均方誤差。PSRN 值越大，圖像之間的像素差異越小。SSIM 分別從亮度、對比度和結構三個方面來評估生成圖像與原圖像的相似性，計算方法如式（12）：

其中μ是圖像的平均值，σ是圖像的方差。C1和C2是用來維持穩定的常數。SSIM 值越大整體結構約相似。L1 loss是輸入和輸出之間絕對差值之和，值越小，表示圖像失真越小。計算方法如式（13）：

表1 統計了本文方法的修復結果與各個方法修復結果幾組數據的平均值，可以看到本文方法在三個評價指標上都優于其他方法。

表1 圖像修復評價指標

4 結語

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