CN111754507A - 一种轻量化的基于强注意力机制的工业缺陷图像分类方法 - Google Patents

Abstract

本发明涉及一种轻量化的基于强注意力机制的工业缺陷图像分类方法，包括以下步骤：1)构建工业缺陷数据集，2)对数据集进行预处理，3)将处理后的数据集进行数据增强操作，4)将增强后的数据集随机分为训练集和验证集，5)使用数据集训练基于注意力机制的分类网络，6)利用训练完的网络模型对工业数据进行分类。本发明采用了较为轻量化的Resnet18网络作为主干架构，网络较为精简，分类网络使用了一种强注意力的约束机制，能够加大对重点区域的关注，能够满足工业缺陷检测在分类精确度和分类速度方面的要求。

Description

一种轻量化的基于强注意力机制的工业缺陷图像分类方法

技术领域

本发明涉及视觉检测技术领域，尤其是一种能够在分类准确率和检测速度上都能达到的工业缺陷分类的检测要求的轻量化的基于强注意力机制的工业缺陷图像分类方法。

背景技术

工业缺陷分类算法是一类适用于在工业流水线上快速地判断出生产的零部件是正常还是异常的一种分类方法。工业缺陷分类算法主要是使用机器视觉算法的算法自动判断出产品等级，从而替代出流水线上的人工成本。由于在工业流水线上需要适合流水线上批量生产的速度，所以需要缺陷分类算法具有较快的分类检测速度。此外，产品的完美程度对于用户体验感是非常重要的，所以工业上对于分类算法的分类精确度要求很高，希望尽量能够检测到缺陷产品。虽然目前有很多机器视觉分类算法应用到工业生产中，但是由于各类分类算法的复杂度和分类精确度不同，造成了在工业曲线分类速度和分类精确度方面还有一些要完善和提高的地方，这也是提高机器视觉算法在工业缺陷应用领域应用体验感和舒适度。

工业缺陷的分类速度快慢取决于网络模型的大小，轻量化精简的网络模型能够获得较快的训练和分类速度，在保证网络分类性能的基础上设计出一个精简的分类网络模型是一项非常重要的工作，也是工业缺陷分类领域一直关注的问题。注意力机制实际上是从人类观察环境的规律而总结出来的，当人类在观察事物时，我们的视觉***会更加倾向于关注图像中辅助判断的部分，而过滤掉一些不相关的部分信息，从而提高判断的准确率，提高模型的收敛速度。注意力机制的优点是使得输出特征图中重要的特征得到加强，而不重要的特征被抑制。最终，可以逐步提升网络的判别能力和收敛速度。

缺陷分类精确度和分类检测速度是工业缺陷检测中关键因素，对缺陷分类的准确度不仅仅可以提高生产产品的质量，而且可以在一定程度上提高工厂的生产效率，而对缺陷产品的分类速度直接影响该算法是否能够在工业生产中使用。产品能够正确分类的关键点是如何设计一种算法来提取到图像的关键特征，进而做出正确的类别判断，以及如何提取一组更为关键有效的特征表示和设计一组较为合理的分类器。

传统图像处理的方式是使用人工设置的方法根据图像的纹理结构来提取图像特征，然后再使用人工神经网络或者机器学习算法来对特征进行分类操作。但是这种传统的分类方法虽然分类速度较快，但是比较依赖于特征提取策略的选择和分类器的选择，有可能还会造成提取的特征与分类器不适合的问题。

目前而言，基于深度学习的分类算法在学术界和工业界得到了较为广泛的研究和应用。但是，对于工业场景中具体的分类问题来说，基于深度学习方法的分类总体精确度表现较好，但是还有提升的空间。基于深度学习的分类网络目前没有能达到很高的分类准确率的原因有很多，但是主要原因有三点：1、在分类性能方面，由于工业业务场景的不同，不同缺陷之间的区分程度不高，造成了缺陷间的误检测出现。2、在分类网络设计方面，由于不同的网络设计方式会捕获到不同的特征信息，但是在网络提取特征时，网络可能没有提取到关键的信息，因此而造成网络分类准确率不高的问题。3、在分类时间方面，由于工业上零部件的缺陷都非常小，因此使用相机采集到的图像都非常大，所以使用复杂的模型往往会在测试推理阶段时间非常长，达不到工业上流水线的速度。

针对工业场景中，工业缺陷图像分类精确度不高的问题。不同缺陷之间的区分程度不高的问题可以尝试从3D图像处理方面进行处理，这样能够获取到输入图像更多的信息，该方法可能从一定程度上，区分缺陷类别间的精度。但是目前基于3D图像处理方面的技术还不是很成熟，使用的不多。

发明内容

本发明要解决的技术问题是：提供一种轻量化的基于强注意力机制的工业缺陷图像分类方法，解决深度神经网络分类过程中对特征定位不准确、分类速度和分类准确率不高的问题。

本发明解决其技术问题所采用的技术方案是：一种轻量化的基于强注意力机制的工业缺陷图像分类方法，包括以下步骤，

S1、构建工业缺陷数据集，数据集由待测零部件的图像数据组成；

S2、对数据集进行处理预处理操作，对图像的缺陷区域进行标注，获得图像的缺陷标注文件；

S3、对处理后的数据集进行数据增强操作，扩充数据集；

S4、将扩充后的数据随机分为训练集和验证集；

S5、利用训练集训练基于注意力机制的分类网络；

S6、利用所训练的分类网络对缺陷图像进行分类。

进一步的说，本发明所述步骤S1中，待检测零部件的图像数据在稳定均匀的光照环境下，通过检测平台上的相机获得。

再进一步的说，本发明所述步骤S2中，预处理操作包括对采集的图像进行去噪、滤波、对比度增强和/或平移翻转；使用label-me标注工具对图像进行标注，获得缺陷位置的标签文件，再将图像的标签文件转换成png图像格式。

再进一步的说，本发明所述步骤S4中，扩增后的数据按照4:1随机分为训练集和验证集。

再进一步的说，本发明所述步骤S5中，分类网络的架构是由7个残差块组成的残差网络架构；每个残差块中有2个卷积核为3x3的卷积操作；整个网络中共有15个卷积核为3x3的卷积操作和4次最大池化下采样操作。

更进一步的说，本发明所述步骤S5中，注意力机制采用强监督方式，计算方式为：H_i,c(x)＝M_i,c(x)*T_i,c(x)；其中，M_i,c(x)表示输入图像的标注信息mask，即权重图像，T_i,c(x)表示网络的特征图，H_i,c(x)是M_i,c(x)和T_i,c(x)这两个向量按照对应位置相乘操作得到的向量。

更进一步的说，本发明所述步骤S5中，分类网络有两个输出，两个输出的损失函数的权重比为3:1。

本发明首先在网络的设计上，采用了较为轻量化的网络架构，以轻量化的Resnet18网络作为主干架构，从而缩短网络的训练和测试推理时间，能够达到工业缺陷检测上对检测速度的要求。此外，本发明中所述的分类网络使用了一种强注意力的约束机制项，在网络训练阶段，强行使得网络关注到关键区域，忽略掉不相关的区域，从而提高网络的分类准确率，加快网络的收敛速度。并且，当训练的数据集较小时，网络容易出现过拟合现象，在这种情况下本发明中注意力机制能够通过原图标注信息的强监督机制，减少网络过拟合情况的发生。此外，本发明的主干网络较为精简能够达到较快的分类速度，网络中的注意力机制还能增加对关键区域的关注力度，使得模型较快地收敛到网络模型的最优点。

本发明的有益效果是：

(1)利用一个轻量化的主干网络架构，通过跳跃连接模块能够更好的利用底层的特征信息，并且能够在一定程度上避免深度神经网络中易出现的过拟合现象；能够满足工业上却分类检测速度的要求，达到实时分类的效果；

(2)利用一种强监督的注意力机制，通过图像的标注信息给分类模型提供一种强监督信息，使得模型在学习的过程中关注到图像中更重要的区域信息，忽略不重要区域的信息，从而解决分类过程中对特征定位不准确的问题，还能使得网络能够更快地达到收敛状态，节省训练时间；

(3)基于强注意力机制能够在分类训练数据量小、网络的学习率设置过大的情况下，强制引导模型关注到图像的重点区域，在一定程度上能够减少模型过拟合情况的发生。

附图说明

图1是本发明的流程图；

图2是本发明的网络架构示意图；

图3是图2中残差块架构示意图；

图4是图2中注意力模块架构示意图。

具体实施方式

现在结合附图和优选实施例对本发明作进一步详细的说明。这些附图均为简化的示意图，仅以示意方式说明本发明的基本结构，因此其仅显示与本发明有关的构成。

如图1-4所示的一种轻量化的基于强注意力机制的工业缺陷图像分类方法，通过一个轻量化的分类网络架构，并利用强注意力机制引导模型更加关注重点区域，从而在保证分类检测速度的情况下提高分类准确率；该方法的分类速度快，分类精度高，能够替代生产车间的人工分选作业，提高工业厂间的流水线的运行速度，降低工厂生产的人力成本；具体包括以下步骤，如图1所示：

步骤S1：数据集构建模块：构建工业缺陷数据集，所述数据需要在工作台上，在稳定均匀的光照环境下，通过检测平台上的相机获得待检测零部件的图像数据。

步骤S2：数据预处理模块，对所述数据集进行处理预处理操作，主要包含对采集的图像进行去噪、滤波、对比度增强和平移翻转等操作，得到质量较好的图像。然后，使用Label-me等标注工具对图像的缺陷区域进行标注，获得图像的缺陷标注文件。

步骤S3：数据增强模块，对缺陷图像进行调整大小操作，获得128x128大小的图像块，然后将图像块进行旋转、翻转、随机噪声等数据增强操作，扩充数据集，将扩充后的数据分为训练集和验证集。

步骤S4：网络训练模块，利用训练集训练所述的基于注意力机制的分类网络。

步骤S5：利用所训练的分类网络对缺陷图像进行分类。

在步骤S1中，数据集的构建之前，更详细的步骤包括：

步骤S101：根据工作台的背景和待检零部件的颜色为相机配置合适的光源，调整相机位置。

在步骤S2中，更详细的数据预处理包括：

步骤S201：对相机采集到的图像进行去噪和滤波等图像操作，获得较为纯净的图像。

步骤S202：使用label-me等标注软件对图像进行标注，获得缺陷位置的标签文件，再将图像的标签文件转换成png图像格式。

在步骤S3中，更详细的数据增强操作包括：

步骤S301：工业上的缺陷数据都是非常大的，具体地，将得到的工业图像调整大小，得到128x128大小的图像块。特别地，对样本标签也做同样的调整大小操作，根据标签类别放入到不同的文件夹中。

步骤S302：将调整大小后得到的图像，通过旋转、翻转、随机噪声和颜色变换等数据增强操作来扩充数据集。当图像在旋转和翻转等改变缺陷位置的操作时，对应图像的标签图像也进行相同的操作。当图像进行随机噪声和颜色变换等缺陷位置不改变的操作时，图像的标签数据不发生改变。最后，将扩增后的数据分为按照4:1随机分为训练集和验证集。

在步骤S4中，更详细的网络训练模块操作包括：

步骤S401：构建基于注意力的生成网络架构，所述网络的架构如图2所示。网络的主干架构是由Resblock组成的残差网络架构，共有7个Resblock，每个Resblock中有2个3x3的卷积操作(如图3)。首先输入128x128x3的图像x，然后经过1个卷积和1个池化层得到64x64x96的张量y1，再对张量有进行1个Resblock和1个最大池化的操作得到32x32x128大小的张量y2，然后再使用张量y2进行2个Resblock和1个最大池化操作得到16x16x256大小的张量y3，再将张量y3经过3个Resblock和1个最大池化操作得到8x8x512大小的张量y4，然后再将张量y4进行一个Resblock的操作得到8x8x512大小的张量y5。将张量y5进行全局平均池化操作(Global Average Pooling)得到维度为512的张量y6，再使用softmax分类函数得到网络的输出output1。如图4所示，张量y5经过1x1的卷积操作降维后再使用sigmoid函数映射得到预测的mask图像输出0，其大小为8x8x1。然后，将输出0在维度方向复制512份，得到512个大小为8x8x1的向量，再使用拼接操作得到了8x8x512大小的张量y7，再将张量y7与张量y5按位相乘，得到了大小为8x8x512的张量y8。再将张量y8进行全局平均池化操作和softmax分类操作得到输出1。

上述网络中共有15个卷积核为3x3的卷积操作和4次最大池化下采样操作，是比较轻量化的网络。网络最后得到两个输出，分别为输出0和输出1。输出0表示经过注意力机制“选择”后的标注信息，其大小为原始图像大小的1/16。输出1表示网络最后的分类结果，即输入图像的类别。

步骤S402：本发明所述的注意力机制原理是使用一种强监督的方式，强行约束网络关注到图像中有差异的部分，能够对细小的缺陷部位产生较强的相应。具体操作计算方式为：

H_i,c(x)＝M_i,c(x)*T_i,c(x)

其中，M_i,c(x)表示输入图像的标注信息mask，即权重图像。对于本发明来说，其大小为8x8，为输入图像对应的标注信息mask经过下采样后得到的，其像素值非0即1。T_i,c(x)表示网络的特征图，其大小也为8x8，是由网络自动学习得到的。然后将这两个向量按照对应位置相乘操作得到了向量H_i,c(x)，即网络只保留M_i,c(x)中值为1的像素值，而M_i,c(x)中值为0的像素值被置零。所以该注意力机制只关注重要区域，将不重要的区域都置零，从而能够加速模型收敛。

步骤S403：所述的注意力机制的工业缺陷分类网络有2个输出，所以就对应着2个损失函数。输出1表示分类网络的输出，使用的是交叉熵损失函数来约束网络。输出0表示的是输入图像的分割图像，大小为原始图像的1/16，若将其上采样到原图大小，即表示的是对原图缺陷的分割结果，所以也采用交叉熵损失函数。本次网络的核心任务是图像分类，所以设定输出1的损失Loss1和输出0的损失函数Loss0的权重比是3:1。分类的损失Loss1的定义如下：

其中，对于分类任务来说，网络的预测值为

表示网络得到的属于各个类别的概率值。y表示该输入图像真实的类别值。输出0的损失函数定义如下：

其中，p表示该像素点的真实类别，

表示该像素点的预测值。网络的总体损失如下：

L＝αLoss1+Loss0

其中α＝3为权重系数，表示对两个损失的加权。

在步骤S5中，利用所训练的分类网络模型对缺陷图像进行分类的详细操作如下：

步骤S501，输入原始图像进行测试任务。将测试图像调整大小为128x128大小的图像块，输入到训练好的分类模型中，输出1即为网络的分类结果。而输出0表示网络生成的输入图像对应的标注图像，在本分类任务中不需要获取。

本发明提出了一个轻量化的分类网络架构，使用轻量化的残差网络为主干网络架构，缩短了缺陷分类时间，能够满足工业生产中对分类算法速度的要求。并且利用强注意力机制来引导网络关注重要的区域，在训练过程中，通过图像的缺陷标注信息来引导网络更加关注到缺陷区域，忽略不重要的区域，从而能够提高分类网络的收敛速度和分类准确率，并且能够从一定程度上减少网络过拟合情况的发生。

本发明将轻量化的残差网络和强注意力机制相结合，在满足工业缺陷图像分类时间要求的情况下，使用注意力机制来提高网络的分类准确率和收敛速度，具有较好的应用价值。

具体表现在：

(a)对于推理速度不高的问题，本发明所涉及到的缺陷分类算法的主干网络是基于残差块的网络，经过实践中轻量化操作，使得该网络更加适合在工业缺陷检测上应用，满足了工业分类检测速度的要求。

(b)对于工业分类任务中对特征的定位不准的问题，本发明采用一种强注意力机制的思想，使得原始的分类网络能够使得网络关注更重要的区域，而忽视不相关的区域，从而提高分类模型的准确率，也加快分类模型的收敛速度。

(c)针对目前分类性能和分类速度的问题，本发明将轻量级的残差网络和注意力机制相结合，从而获得总体性能上的突破。

以上说明书中描述的只是本发明的具体实施方式，各种举例说明不对本发明的实质内容构成限制，所属技术领域的普通技术人员在阅读了说明书后可以对以前所述的具体实施方式做修改或变形，而不背离发明的实质和范围。

Claims (7)

1.一种轻量化的基于强注意力机制的工业缺陷图像分类方法，其特征在于：包括以下步骤，

S1、构建工业缺陷数据集，数据集由待测零部件的图像数据组成；

S2、对数据集进行处理预处理操作，对图像的缺陷区域进行标注，获得图像的缺陷标注文件；

S3、对处理后的数据集进行数据增强操作，扩充数据集；

S4、将扩充后的数据随机分为训练集和验证集；

S5、利用训练集训练基于注意力机制的分类网络；

S6、利用所训练的分类网络对缺陷图像进行分类。

2.如权利要求1所述的一种轻量化的基于强注意力机制的工业缺陷图像分类方法，其特征在于：所述步骤S1中，待检测零部件的图像数据在稳定均匀的光照环境下，通过检测平台上的相机获得。

3.如权利要求1所述的一种轻量化的基于强注意力机制的工业缺陷图像分类方法，其特征在于：所述步骤S2中，预处理操作包括对采集的图像进行去噪、滤波、对比度增强和/或平移翻转；使用label-me标注工具对图像进行标注，获得缺陷位置的标签文件，再将图像的标签文件转换成png图像格式。

4.如权利要求1所述的一种轻量化的基于强注意力机制的工业缺陷图像分类方法，其特征在于：所述步骤S4中，扩增后的数据按照4:1随机分为训练集和验证集。

5.如权利要求1所述的一种轻量化的基于强注意力机制的工业缺陷图像分类方法，其特征在于：所述步骤S5中，分类网络的架构是由7个残差块组成的残差网络架构；每个残差块中有2个卷积核为3x3的卷积操作；整个网络中共有15个卷积核为3x3的卷积操作和4次最大池化下采样操作。

6.如权利要求5所述的一种轻量化的基于强注意力机制的工业缺陷图像分类方法，其特征在于：所述步骤S5中，注意力机制采用强监督方式，计算方式为：H_i,c(x)＝M_i,c(x)*T_i,c(x)；其中，M_i,c(x)表示输入图像的标注信息mask，即权重图像，T_i,c(x)表示网络的特征图，H_i,c(x)是M_i,c(x)和T_i,c(x)这两个向量按照对应位置相乘操作得到的向量。

7.如权利要求6所述的一种轻量化的基于强注意力机制的工业缺陷图像分类方法，其特征在于：所述步骤S5中，分类网络有两个输出，两个输出的损失函数的权重比为3:1。

Images