許國良 毛 驕

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

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

1 引言

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

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

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

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

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

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

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

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

2 相關(guān)內(nèi)容

2.1 語義分割

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

2.2 小樣本學習

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

2.3 小樣本分割

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

3 基于協(xié)同注意力的小樣本手機屏幕缺陷分割網(wǎng)絡(luò)

3.1 網(wǎng)絡(luò)提出

3.2 網(wǎng)絡(luò)架構(gòu)

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

圖1 基于協(xié)同注意力的小樣本手機屏幕缺陷分割網(wǎng)絡(luò)的網(wǎng)絡(luò)架構(gòu)圖

3.2.1 缺陷圖像特征提取模塊

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

圖2 特征提取模塊示意圖

圖3 交叉注意力塊示意圖

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

3.2.2 特征增強模塊

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

圖4 特征增強模塊架構(gòu)圖

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

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

3.2.3 解碼器

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

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

4 實驗與分析

4.1 手機屏幕缺陷圖像數(shù)據(jù)集

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

表1 手機屏幕缺陷圖像數(shù)據(jù)集

4.2 實驗細節(jié)

本文的全部實驗均在PyTorch深度學習框架下實現(xiàn)，使用了動量為0.7的SGD優(yōu)化器來優(yōu)化所有的網(wǎng)絡(luò)模型，初始學習率設(shè)置為0.01，每訓練10000次，學習率降低一半，網(wǎng)絡(luò)總共訓練100000次。最后采用聯(lián)合的交叉熵損失函數(shù)來完成網(wǎng)絡(luò)的訓練。

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

4.3 實驗結(jié)果及分析

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

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

圖6 不同分割網(wǎng)絡(luò)對手機屏幕缺陷圖像的分割效果可視化

表2 不同分割網(wǎng)絡(luò)模型在手機屏幕缺陷數(shù)據(jù)集的性能比較

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

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

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

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

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

5 結(jié)束語

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