基于協同注意力的小樣本的手機屏幕缺陷分割

許國良 毛 驕

①(重慶郵電大學通信與信息工程學院 重慶 400065)

②(重慶郵電大學電子信息與網絡工程研究院 重慶 400065)

1 引言

手機屏幕的生產過程極為復雜，在生產的過程中容易出現各種各樣的缺陷，因此要保證手機屏幕高質量的要求，對其進行質量檢測是生產過程中必不可少的一道工序。人工檢測容易受到主/客觀因素的影響，使得檢測速度慢、效率低、漏檢率高，不適應手機屏幕的快速生產模式。基于機器學習的缺陷檢測方法很大程度上依賴缺陷特征提取算法，針對特定缺陷的檢測效果較好，然而并不適用于種類繁多復雜的手機屏幕缺陷的檢測。近年來，語義分割網絡的出現不僅克服了傳統檢測方法特征提取的問題，還提供缺陷的定位和分類，在缺陷檢測領域廣泛應用。Tabernik等人[1]使用雙分支網絡來實現圖像的缺陷檢測，但輸出的分割圖像分辨率為原始輸入圖像的1/8，這會丟失尺寸較小缺陷的相關信息。一些缺陷分割方法[2,3]利用多階段的全卷積網絡對輸入圖像的粗略預測與精細分割，以完成較好的分割效果。另外，一系列基于無監督的缺陷檢測方法[4–7]通過正樣本完成對重構網絡的訓練，然后將輸入樣本與修復后的圖像進行比較，以得到輸入樣本的分割區域，這類方法很好地解決了大量標注樣本的需求問題。不同的是，文獻[4]的重構網絡是一個卷積去噪自編碼器，文獻[5]的重構網絡則結合了生成對抗網絡(Generative Adversarial Network,GAN)和自動編碼器，文獻[6]融合多個全卷積自編碼器重建紋理圖像。而文獻[7]在深度卷積生成對抗網絡(Deep Convolution Generative Adversarial Networks, DCGAN)的基礎上，引入一個額外的編碼器，以此完成對紋理圖像的重建。

然而，上述的分割網絡需要大規模的缺陷圖像樣本來訓練網絡以獲取較好的分割效果。但隨著手機屏幕生產技術的提高，能采集到的缺陷類別、數量都是有限的，構建的手機屏幕缺陷圖像數據集不足以完成大數據驅動的分割網絡的訓練。因此，如何利用少量的手機屏幕缺陷圖像來很好地完成手機屏幕缺陷分割成為亟待解決的問題。

在很多場景下，收集大量的有標簽的數據是非常昂貴、困難甚至不可能的。受這一事實的影響，小樣本學習就應運而生，這種方法能夠利用先驗知識，通過泛化和類比從少量樣本中獲取知識，同時能快速適應于新任務中。而小樣本分割網絡旨在利用少量的具有分割標簽的樣本圖像來指導測試圖像的分割過程，以此獲得測試圖像的分割結果。目前，劉宇軒等人[8]提出了一個注意力譜生成器來融合全局相似性和局部相似性以實現小樣本分割。董陽等人[9]提出基于U-net的原型網絡模型完成MR腦瘤圖像的分割。羅善威等人[10]構建孿生殘差網絡，通過融合空域相似度和頻域相似度獲得分割結果。

針對手機屏幕缺陷，本文設計了一種基于協同注意力的小樣本手機屏幕缺陷分割網絡(Co-Attention Segmentation Network, Co-ASNet)。整個網絡采用編碼器-解碼器的架構，在編碼階段引入交叉注意力塊(Criss-Cross Attention Block, CC-block)來提取不同尺寸缺陷的特征信息，同時引入的協同注意力更多地用來交換支持圖像和查詢圖像的特征信息以減少類內差距，以此來增強缺陷的特征表示。此外，本文采用了支持圖像和查詢圖像聯合損失函數來訓練網絡，以此來獲得更好的分割效果。本文的貢獻如下：

(1)設計一種基于協同注意力的小樣本手機屏幕缺陷分割網絡，用于手機屏幕缺陷分割。在編碼模塊引入交叉注意力模塊來學習上下文信息以更好地表征不同尺寸、不同位置的缺陷信息；

(2)本文引入了協同注意力來增強支持圖像與查詢圖像之間的信息交互，使支持圖像更好地指導查詢圖像的缺陷分割，同時利用支持圖像和查詢圖像的聯合損失函數來提升缺陷分割效果；

(3)本文網絡的訓練采用的是小樣本學習中的C-wayK-shot(C>1)的訓練策略，而不是其他小樣本分割網絡所使用的對于每一類的1-wayK-shot策略。

針對手機屏幕缺陷圖像數據集，本文所提出的小樣本分割網絡模型可以取得較好的缺陷分割效果，相較于其他分割網絡，該模型能夠使用較少的缺陷樣本完成較好的缺陷分割。

2 相關內容

2.1 語義分割

語義分割網絡是在目標檢測網絡的基礎上，對輸入圖像的每個像素進行分類，標注出圖像中每個像素的所屬的對象類別。2014年，Shelhamer等人[11]提出的全卷積網絡(Fully Convolutional Networks,FCN)去掉了卷積神經網絡的全連接層，用卷積層代替得到分割圖像。隨后，Ronneberger等人[12]在全卷積網絡的基礎上設計了U-net，該網絡采用編碼器-解碼器的架構完成分割。SegNet[13]使用池化索引，即記錄下最大像素值的位置，以便于解碼時的上采樣。DeepLab系列[14–17]則是引入了空洞卷積來增大感受野，以捕獲不同尺度特征的上下文信息。然而，以上分割網絡還是依賴大量標簽數據來完成分割，針對少量的標簽數據，其分割效果較差。本文在少量的手機屏幕缺陷圖像的基礎上，完成屏幕缺陷的分割。

2.2 小樣本學習

小樣本學習是一種利用少量有監督信息的機器學習問題，其重點在于在標簽數據有限的情況下，網絡能夠快速地進行學習，并能夠泛化到其他新任務中。目前，小樣本學習大致分為3類：一是遷移學習[18]，其將預訓練模型遷移到新的模型中，通過少量樣本使得模型適應新的應用場景。二是元學習[19,20]，其模型利用以往的知識經驗來指導新任務的學習，具有學會學習的能力。三是基于度量學習的方法[21–23]，通過學習樣本間的距離來減少同類樣本距離，增大異類樣本距離。孿生網絡利用兩個并行網絡來衡量輸入的成對圖像的相似程度。匹配網絡提出一種小樣本學習的通用模型，通過將支持圖像和查詢圖像映射到同一個嵌入空間，利用余弦相似度度量兩者的相關性以完成分類。而原型網絡在匹配網絡的基礎上，學習每一個類別的原型表示以計算測試圖像與每一類別的相似性。關系網絡則是基于小樣本學習的網絡框架，將相似性度量模型化，利用網絡學習兩者的相關性。現實中，許多場景樣本采集的數量有限或者標注數據的成本太大，針對這些問題，小樣本學習能夠很好地解決。

2.3 小樣本分割

小樣本分割主要利用少量標簽樣本來訓練模型，使得模型能泛化新的分割任務。文獻[24]和文獻[25]利用條件分支對查詢圖像進行加權生成分割圖像。相似性引導網絡(Similarity Guidance network for ONE-shot semantic segmentation, SG-ONE)[26]提出了掩膜平均池化來提取支持圖像的表征向量，提高分割效果。文獻[27]利用了原型計算相似度，文獻[28]則是引入了迭代優化的方法改進模型，以迭代方式優化分割效果。Nguyen等人[29]改進了SG-One網絡。Liu等人[30]提出了用于小樣本分割的交叉參考網絡(Cross-Reference Network, CRNet)，該網絡利用交叉參考機制可以更好地在兩幅圖像中找到相同的目標。為了更好地實現在少量樣本下的手機屏幕缺陷分割效果，本文采用小樣本分割網絡框架進行改進。

3 基于協同注意力的小樣本手機屏幕缺陷分割網絡

3.1 網絡提出

3.2 網絡架構

針對手機屏幕缺陷，本文提出了一種基于協同注意力的小樣本手機屏幕缺陷分割網絡(Co-ASNet)用于手機屏幕缺陷分割。與其他小樣本分割網絡架構相似，所提出的Co-ASNet為編碼器-解碼器的架構，不同的是編碼器是由特征提取模塊和特征增強模塊構成的，如圖1所示。在特征提取模塊中，為了能夠更好地提取到不同尺寸缺陷的特征信息，本文在特征提取模塊引入交叉注意力塊來有效地獲取缺陷圖像的上下文信息；而在特征增強模塊中，本文沿用了SG-One網絡中利用掩碼平均池化生成支持圖像的代表特征，通過余弦相似度量來指導查詢圖像的分割過程，同時引入協同注意力來加強查詢圖像與支持圖像之間的信息交互，增強兩者在分割過程中的相關性，以此更好地完成查詢圖像的缺陷分割。在解碼器部分，本文采用了U-net的反卷積部分，最后利用聯合的支持/查詢圖像損失函數來訓練網絡。另外，虛線部分只在訓練過程中使用，測試過程只使用實線部分。

圖1 基于協同注意力的小樣本手機屏幕缺陷分割網絡的網絡架構圖

3.2.1 缺陷圖像特征提取模塊

針對手機屏幕缺陷的大小、位置靈活等特點，本文在特征提取模塊中引入交叉注意力塊來獲取缺陷圖像的上下文信息以有效獲取不同尺寸缺陷的特征信息，增強缺陷圖像的特征表示。特征提取模塊如圖2所示，其包含了5個卷積塊和兩個交叉注意力塊，每個卷積塊包含一個3×3的卷積層，一個歸一化層，一個ReLU激活層以及一個2×2的最大池化層。

圖2 特征提取模塊示意圖

圖3 交叉注意力塊示意圖

雖然每一個交叉注意力塊可以聚合某個位置上的水平和垂直方向上的特征信息，但這個位置周圍其他不在同一行/列的位置之間的相關性無法獲得。而文獻[31]連續采用兩個交叉注意力塊可以從所有像素中獲取全圖像的上下文信息，生成具有密集和豐富的上下文信息的新特征，在很大程度上節約了時間和GPU內存。本文采用兩個交叉注意力塊來獲取手機屏幕缺陷的豐富的特征信息。

3.2.2 特征增強模塊

在小樣本分割網絡中，為了使查詢圖像的分割效果更好，通常使用支持圖像及其分割掩膜圖像對查詢圖像進行條件處理，指導完成查詢圖像分割過程。本文利用SG-One網絡所提出掩膜平均池化來獲取支持圖像的代表特征，并利用相似度度量的方式來建立支持圖像特征和查詢圖像特征之間的關系；同時為了使查詢圖像與支持圖像的信息交互得更加充分，本文還引入了協同注意力來強調查詢圖像與支持圖像的相同目標的相關性。圖4所展示的是特征增強模塊的整體架構。

圖4 特征增強模塊架構圖

(b)K-shot。K>1，即對于每一類缺陷，支持集中包含著K個缺陷樣本。本文將這K個支持樣本的特征圖進行加權平均來作為這類缺陷的特征圖，同時對這K個缺陷樣本對應的掩膜圖像也進行加權平均得到一個新的掩膜圖像，使用新特征圖和新的掩膜圖像來完成掩膜平均池化過程。

(2)協同注意力。在掩膜平均池化時，雖然利用支持圖像及其掩膜圖像生成的代表特征來進一步指導查詢圖像的特征的形成，但兩者信息的交互是單方面從支持圖像對查詢圖像。但在小樣本分割網絡訓練過程中，缺陷圖像在不同的訓練任務中扮演著不同的角色，可能作為支持圖像，可能作為查詢圖像，這說明在訓練時，支持集和測試集是可以相互影響的。為了進一步加強支持集與查詢集的相同目標之間的信息交互，本文引入協同分割的機制來對查詢缺陷圖像特征進行細化，如圖4所示。

3.2.3 解碼器

類似U-net，本文在解碼的過程中，采用跳躍連接的方式對提取到的特征生成缺陷分割圖像，如圖5所示。本文將特征提取模塊中的第3個卷積層和第4個卷積層與解碼器中第一個反卷積層的輸出和第2個卷積層的輸出相加，共同作為下一層的輸入，最終解碼器輸出的分割圖像與掩膜圖像一樣大。在本文中，由于在特征加強模塊中引入了協同分割機制，這不管是對支持圖像的特征還是對查詢圖像的特征都會有更加豐富的特征表示，所以在訓練階段，訓練網絡的損失的函數也做一定的改進。不管是支持圖像還是查詢圖像，對兩者的特征進行恢復得到的分割圖像都會有利于網絡的訓練，增強缺陷分割的效果。本文聯合了支持圖像和查詢圖像的交叉熵損失函數，共同訓練網絡，如式(10)所示。

圖5 編碼-解碼過程示意圖

4 實驗與分析

4.1 手機屏幕缺陷圖像數據集

本文是針對手機屏幕缺陷所提出的小樣本分割方法，為了驗證所提出的網絡的性能，構建手機屏幕缺陷數據集，如表1所示。整個數據集從手機屏幕工廠生產線采集到的3750×2098高分辨率手機屏幕缺陷圖像通過剪切成160×160大小的缺陷圖像構成。其中，960張缺陷圖像用于整個網絡的訓練，140張缺陷圖像用于測試。由于構建的手機屏幕缺陷數據集的數據量較少，在對缺陷進行分類時，缺陷分成了3個類別：點缺陷、線缺陷以及面缺陷，以此來完成手機屏幕缺陷分割。

表1 手機屏幕缺陷圖像數據集

4.2 實驗細節

本文的全部實驗均在PyTorch深度學習框架下實現，使用了動量為0.7的SGD優化器來優化所有的網絡模型，初始學習率設置為0.01，每訓練10000次，學習率降低一半，網絡總共訓練100000次。最后采用聯合的交叉熵損失函數來完成網絡的訓練。

在訓練時，跟之前小樣本分割網絡所采用的策略不同，它們的每一次訓練任務只有一種類型的K個樣本參與訓練，相當于1-wayK-shot的訓練策略。本文采用了通過小樣本學習類似的訓練策略C-wayK-shot的方式來訓練本文構建的Co-ASNet。具體地，在實驗過程中，設置了3-way 1-shot和3-way 5-shot兩種訓練模式，也就是說，在每一次的訓練任務中，參與訓練的缺陷類別包括了點、線、面3類，每一類缺陷的支持樣本分別為1和5，以此來完成網絡的整體訓練。在測試時，采用了與訓練相同的策略，但不同的是，在訓練時，不論是支持樣本還是查詢樣本都是從訓練集中選取的，而測試過程是從訓練集中獲取支持集，從測試集中抽取測試樣本，這些測試樣本是在訓練過程中未被網絡訓練過的樣本。

4.3 實驗結果及分析

為了評估網絡對手機屏幕缺陷檢測的效果，使用的評判指標是像素準確率(Pixel Accuracy, PA)、平均像素準確率(Mean Pixel Accuracy, MPA)、平均交并比(Mean Intersection over Union, MIoU)以及頻率加權交并比(Frequency Weighted Intersection over Union, FWIoU)，具體的定義如式(11)—式(14)。

在本文的實驗驗證過程中，除了實現本文所提出的網絡，還對比了傳統的語義分割網絡U-net以及小樣本分割網絡SG-One，具體的實驗結果如表2所示，可以看出，本文所提出的網絡在評判指標上均取得比較好的結果。具體來說，U-net對于只有少量的手機屏幕缺陷樣本來說，其分割的效果不是很理想，而SG-One網絡不論是在1-shot還是在5-shot, MIoU值相較于U-net有所提高。當K=5時，本文所提出的網絡在SG-One的基礎上，MPA和MIoU值達到了最大，MPA為0.6711, MIoU為0.5771，對手機屏幕缺陷語義分割的效果進一步提高。另外，圖6也給出了各種網絡對手機屏幕缺陷的分割效果圖，在測試1-shot時，本文使用表1的缺陷圖像作為支持集，在5-shot測試時，支持集也包含表1的缺陷圖像。

圖6 不同分割網絡對手機屏幕缺陷圖像的分割效果可視化

表2 不同分割網絡模型在手機屏幕缺陷數據集的性能比較

除此之外，對于本文所提出的網絡自身，引入了協同分割機制來增強缺陷特征的表達，也改進了損失函數來提升網絡的性能。為了能夠更好地展示引入的交叉注意力、協同分割機制和改進的損失函數對手機屏幕缺陷語義分割的影響，本文也做了相關消融實驗，如表3所示。在只將支持圖像的注意力特征圖來細化查詢圖像特征信息，這雖然在SGOne的基礎上分割的效果有一定的提升，但對于支持圖像特征與查詢圖像特征之間的相互交互來講，分割的效果有待提高。而當改進損失函數，聯合查詢圖像和支持圖像一起來訓練，加強了兩者相同目標之間的信息的交互，使得支持圖像更好地指導查詢圖像的分割過程，提升了分割效果。不管是在1-shot，還是在5-shot的設置下，本文所提出的網絡取得的分割效果更好。在1-shot下，MIoU值達到0.5588，而在5-shot設置下，MIoU值達到最大0.5771。

表3 在手機屏幕缺陷圖像數據集上的分割結果(MIoU)

最后，對于Co-ASNet所引入的交叉注意力以及聯合訓練的協同注意力機制在分割結果上帶來的效果上的提升，本文在圖7、圖8分別針對1-shot和5-shot的設置，不同的消融實驗對手機屏幕缺陷分割效果的可視化展示。

圖7 1-shot下的手機屏幕缺陷圖像的分割效果可視化

圖8 5-shot下的手機屏幕缺陷圖像的分割效果可視化

5 結束語

本文針對手機屏幕缺陷，提出了基于協同注意力的小樣本的手機屏幕缺陷分割網絡。整個網絡由特征提取模塊、特征增強模塊以及解碼器3個部分構成。在特征提取模塊中，引入交叉注意力模塊以獲取更加豐富的特征信息；其次，為了加強支持圖像與查詢圖像特征的相同目標的特征信息交互，在特征增強模塊采用了協同注意力來處理；最后在解碼器中，逐步生成缺陷圖像的分割圖像。在網絡的訓練過程中，聯合支持圖像與查詢圖像損失函數共同訓練網絡，增強網絡分割效果。雖然所提出的網絡緩解了工業上手機屏幕缺陷樣本較少的問題，但在缺陷分割場景中的分割效果還達不到工業化缺陷檢測的標準，對分割的效果還有待提升。因此，未來的研究還要著重提升網絡的分割效果，以實現精確的定位與分類。