基于协同注意力的小样本的手机屏幕缺陷分割

许国良 毛 骄

①(重庆邮电大学通信与信息工程学院 重庆 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个部分构成。在特征提取模块中，引入交叉注意力模块以获取更加丰富的特征信息；其次，为了加强支持图像与查询图像特征的相同目标的特征信息交互，在特征增强模块采用了协同注意力来处理；最后在解码器中，逐步生成缺陷图像的分割图像。在网络的训练过程中，联合支持图像与查询图像损失函数共同训练网络，增强网络分割效果。虽然所提出的网络缓解了工业上手机屏幕缺陷样本较少的问题，但在缺陷分割场景中的分割效果还达不到工业化缺陷检测的标准，对分割的效果还有待提升。因此，未来的研究还要着重提升网络的分割效果，以实现精确的定位与分类。