一种基于MSTA-YOLOv5的轻量化齿轮表面缺陷检测方法

技术领域

本发明涉及产品缺陷检测技术领域，具体涉及一种基于MSTA-YOLOv5的轻量化齿轮表面缺陷检测方法。

背景技术

随着科学技术的发展与社会需求变化，大型多面复杂结构工件在工业生产中更加普及，齿轮是机械工业中应用广泛的传动部件，其质量在生产中尤为重要。但在实际生产过程中，由于工艺流程、生产设备和现场环境等因素的影响，造成齿轮表面出现各种缺陷，如果不及时处理，将会影响齿轮表观质量、性能和使用寿命，导致企业生产效益下降。因此，需要对齿轮表面进行检测，而传统的人工检测工作量大，易造成检测人员视觉疲劳，出现漏检、错检。

近年来，随着机器视觉技术快速发展，基于机器视觉的检测技术被应用在产品表面质量检测中。而现在的齿轮缺陷检测技术大部分采用数字化图像处理技术，但这种技术的处理方式、算法单一，面对表面复杂性较高的齿轮来说，很难有效的提取缺陷目标，导致检测效果不理想。

如公开号为CN115187820A的专利文件公开了轻量化的目标检测方法、装置、设备、存储介质，在YOLOv4网络结构中采用了ShuffleNetv作为特征提取模块，但是参数量和计算量非常大，且特征提取模块中用到了SE注意力机制，具有精度不足的缺点；

如公开号为CN112990325A的专利文件公开了一种面向嵌入式实时视觉目标检测的轻型网络构建方法，用到了CBAM注意力机制，具有轻量化的优点，但是随着轻量化提高其精度损失较大，且该技术方案中采用了Focus切片操作，增加了参数量，减弱了轻量化的优势；

如公开号为CN114898171A的专利文件公开了一种适用于嵌入式平台的实时目标检测方法，虽然达到了轻量化的效果，但是精度还是造成了比较大的损失。

随着人工智能技术的发展，深度学习方法因其在处理背景复杂、缺陷微弱的工业图像时展现出卓越的性能优势，被广泛应用于图像处理和工件质量检测领域。采用深度学习方法，能够准确地对齿轮表面缺陷进行语义识别和分割，降低了背景及其他因素的干扰，从而有效提升了检测准确性。虽然有大量研究对不同的目标检测网络进行改进并对工业产品进行缺陷检测，达到了可观的效果，但没有针对企业内需要体积小、计算参数更少的模型进行研究，这样的模型在成本预算低、算力相对不足的设备中也能达到不错的检测速度和准确率。

发明内容

深度学习方法在图像分类领域已经极大地提高了准确度，但目前基于深度学习的目标检测算法计算资源需求太大、内存消耗严重使得成本较高。本发明针对上述问题以及企业需要低延迟模型并且移动设备终端需要既快又准确的小模型问题，提供了一种基于MSTA-YOLOv5的轻量化齿轮表面缺陷检测方法，实现齿轮表面缺陷的检测与自动分拣，能够实现齿轮表面缺陷检测的检测效率提高。

本发明解决技术问题的技术方案为：

一种基于MSTA-YOLOv5的轻量化齿轮表面缺陷检测方法，包括以下步骤：首先获取齿轮表面缺陷图像，并对所述的图像进行标注和划分，构建齿轮表面缺陷数据集；然后构建MSTA-YOLOv5检测模型，基于齿轮表面缺陷数据集对MSTA-YOLOv5检测模型进行训练；最后将待检测的齿轮缺陷图像送入训练好的MSTA-YOLOv5检测模型，获取检测齿轮的缺陷类型。

所述的MSTA-YOLOv5检测模型的包括：

输入部分：将齿轮表面缺陷图像输入MSTA-YOLOv5网络，进行自适应锚框计算和Mosaic9数据增强；

骨干部分：特征提取主干网络采用ShuffleNetv2架构，包括依次连接的CBRM操作、第一下采样层、第二卷积归一化层、第二下采样层、第三卷积归一化层、第三下采样层、第四卷积归一化层；经过下采样层处理后的齿轮表面缺陷图像利用1*1卷积进行特征提取之后得到的3个齿轮表面缺陷特征图分别记为S2,S3,S4；

颈部部分：颈部Neck结构采用FPN+PAN，FPN层自顶向下传递强语义信息，S4经过3*3卷积，得到特征图记为Q4，Q4经过转置卷积上采样后与S3相连接，再经过3*3卷积，得到特征图记为Q3；Q3经过转置卷积上采样后与S2相连接，再经过3*3的卷积，得到的特征图记为Q2；

PAN自底向上传递强定位信息，特征图Q2作为底层特征R2，R2经过下采样后，与Q3相连接，得到特征图记为R3；R3经过下采样后与Q4相连接，得到的特征图记为R4；R2、R3、R4分别经过3*3的卷积，得到特征图T2、T3、T4；

在颈部Neck结构的后3个C3模块之后分别集成一个AMECA注意力模块，分别将特征图T2、T3、T4作为原始输入特征图，分别经过全局平均池化模块与全局最大池化模块，并将得到的两个特征图进行相加，压缩空间信息，随后使用1*1卷积学习通道注意力信息，得到的通道注意力信息与原始输入特征图结合，最终得到具体的通道注意力特征图D1、D2、D3；

输出部分：将特征图D1、D2、D3分别输入YOLOv5-MSTA检测头网络，最终得到检测结果。

进一步地，所述的Mosaic9数据增强包括：首先从总数据集中取出一个批量的数据，每次从中随机取出9张图片，进行随机位置的裁剪、缩放，合成新图片；上述过程重复batch-size次，最后得到包括batch-size个经过了Mosaic9数据增强后图片的一个批量的新数据，再传递给神经网络进行训练。

进一步地，所述的CBRM操作包括Conv、BN、ReLU和MaxPool。

进一步地，所述的第一下采样层、第二下采样层、第三下采样层中均包括Shuffle_Block(d)模块，所述的Shuffle_Block(d)将输入特征分为两个分支，左侧分支有2个卷积层，分别为步长为2的3×3深度卷积和1×1的普通卷积；右侧分支上有三个卷积层，分别为1×1的普通卷积、步长为2的3×3深度卷积和1×1的普通卷积；左右两侧分支通过Concat进行拼接来使左右两侧的特征进行融合，最后通过通道混洗操作来启用两个分支之间的信息通信。

进一步地，所述的第二卷积归一化层、第三卷积归一化层、第四卷积归一化层中均包括Shuffle_Block(c)模块，所述的Shuffle_Block(c)模块对每个通道进行分流，划分为两个分支，根据要减少模型的碎片化程度准则，在左边的分支上没有进行任何操作，在右边的分支上有3个卷积层，分别是1×1普通卷积、3×3深度卷积、1×1普通卷积，1×1普通卷积、3×3深度卷积、1×1普通卷积三个卷积层拥有相同的输入和输出通道，其中两个所述的1×1Conv不再是组卷积而改变为普通的卷积，在所述的3个卷积之后，两个分支是通过Concat进行拼接。

进一步地，所述的转置卷积上采样的运算步骤包括：

(1)在输入特征图元素间填充s-1行、列0，其中s表示转置卷积的步距；

(2)在输入特征图四周填充k-p-1行、列0，其中k表示转置卷积的kernel_size大小，p为转置卷积的填充；

(3)将卷积核参数上下、左右翻转；

(4)做正常卷积运算，填充0，步距1。

进一步地，所述的AMECA注意力模块的流程包括：

(1)首先输入特征图X，特征图X的维度是H*W*C；

(2)对输入特征图X进行空间特征压缩；在空间维度，使用全局平均池化GAP，得到1*1*C的特征图F1；使用全局最大值池化GMP，得到1*1*C的特征图F2；

(3)将F1与F2融合得到1*1*C的特征图F3，获得更高层次的语义信息；

(4)对融合后的特征图F3进行通道特征学习；通过1*1*1卷积，学习不同通道之间的重要性，此时输出的特征图F4维度还是1*1*C；

(5)将特征图F4经过σ函数得到F41；

(6)最后是通道注意力结合，将通道注意力的特征图F41与原始输入特征图X，进行逐通道乘积，最终输出具有通道注意力的特征图X'；

其中，H，W和C分别表示输入特征图的高、宽和通道数，σ代表激活函数；

所述的特征图T2、T3、T4作为输入特征图X，分别得到的输出特征图X'为特征图D1、D2、D3。

一种计算机可读介质其上存储有计算机程序用于执行如上所述的方法。

本发明的有益效果：

本发明通过在输入端采用Mosaic9数据增强，在丰富数据集的同时增加了小样本目标，提升网络的训练速度和泛化能力；为了便于模型部署，本发明使用ShuffleNetv2作为骨干网络提取特征，通道重排实现跨组信息交流，构建YOLOv5轻量级神经网络模型，降低网络参数量的同时提升模型检测速度；通过采用转置卷积方式进行上采样，实现了语义级别的上采样，使特征含有更强的语义信息的同时能够使网络更加轻量化；最后在Neck结构中加入AMECA注意力机制，通过注意力模块调整模型通道的信息提取方式，加强通道特征，使得缺陷检测更加精准，进一步增强齿轮缺陷特征提取能力，提高齿轮缺陷模型检测的性能。

附图说明

图1为本发明的MSTA-YOLOv5检测模型的网络结构图；

图2为本发明的Mosaic9数据增强的流程结构图；

图3为本发明的ShuffleNetv2的两种模块的结构图；

图4为本发明的AMECA注意力模块的结构图；

图5为YOLOv5的网络结构图；

具体实施方式

为了能清楚说明本方案的技术特点，下面通过具体实施方式，并结合其附图，对本发明进行详细阐述。下文的公开提供了许多不同的实施例或例子用来实现本发明的不同结构。为了简化本发明的公开，下文中对特定例子的部件和设置进行描述。此外，本发明可以在不同例子中重复参考数字和/或字母。这种重复是为了简化和清楚的目的，其本身不指示所讨论各种实施例和/或设置之间的关系。应当注意，在附图中所图示的部件不一定按比例绘制。本发明省略了对公知组件和处理技术及工艺的描述以避免不必要地限制本发明。

一种基于MSTA-YOLOv5的轻量化齿轮表面缺陷检测方法，包括以下步骤：首先获取齿轮表面缺陷图像，并对所述的图像进行标注和划分，构建齿轮表面缺陷数据集；然后构建MSTA-YOLOv5检测模型，基于齿轮表面缺陷数据集对MSTA-YOLOv5检测模型进行训练；最后将待检测的齿轮缺陷图像送入训练好的MSTA-YOLOv5检测模型，获取检测齿轮的缺陷类型。本发明齿轮表面缺陷类型包括齿底黑皮、齿面黑皮和磕碰三种类型。

如图1所示，所述的MSTA-YOLOv5检测模型的包括：

输入部分：将齿轮表面缺陷图像输入MSTA-YOLOv5网络，进行自适应锚框计算和Mosaic9数据增强；

骨干部分：特征提取主干网络采用ShuffleNetv2架构，包括依次连接的CBRM操作、第一下采样层、第二卷积归一化层、第二下采样层、第三卷积归一化层、第三下采样层、第四卷积归一化层；经过下采样层处理后的齿轮表面缺陷图像利用1*1卷积进行特征提取之后得到的3个齿轮表面缺陷特征图分别记为S2,S3,S4；

如图5、图1所示，本发明在传统的YOLOv5模型基础上，用ShuffleNetV2架构替代CSPDarknet53作为特征提取网络，构建YOLOv5轻量级神经网络模型。ShuffleNetV2既继承了ShuffleNet分组卷积和通道重排的特点，又遵循了设计轻量级网络的4条准则。在同等条件下，ShuffleNetV2相比其他模型速度快，而且准确度更好。MSTA-YOLOv5模型输入由ShuffleNetV2提取的目标特征量，根据每次迭代返回的损失值，自适应的调整网络模型的参数，等待损失值收敛趋于稳定，便可以得到评估指标最佳的检测模型。大大减少了模型的参数量Parameters和计算量FLOPs，减小了模型的大小。CSPDarknet53和ShuffleNetv2两个架构的层数、参数量和计算量的对比如表1所示。

表1 主干网络参数量对比

其中FLOPs就是指计算量的大小，对于卷积层而言，FLOPs的计算公式如下：

FLOPs＝2HW(C

其中，C

FLOPs＝HW(C

对于卷积层而言，参数量Parameters的计算公式如下：

parameters＝Co×(Ci×K×K+1) (3)

其中，Co是输出通道数，Ci是输入通道数，K指的是卷积核的大小；

H、W分别表示输入特征图的高、宽；

如图3所示，ShuffleNetv2主干有两种模块分别为Shuffle_Block(d)、Shuffle_Block(c)，Shuffle_Block(c)和Shuffle_Block(d)在yaml配置文件里通过步长就可以区分，Shuffle_Block(c)的步长stride＝2，Shuffle_Block(d)的步长stride＝1，这两个在本发明中是交替使用的。将原始YOLOv5网络输入端的Focus切片替换成CBRM，一个3*3的卷积，将主干网络所有Conv+C3替换为Shuffle_Block，去掉了SPP和后面的一个C3结构，因为SPP的并行操作会影响速度。

颈部部分：颈部Neck结构采用FPN+PAN，FPN层自顶向下传递强语义信息，S4经过3*3卷积，得到特征图记为Q4，Q4经过转置卷积上采样后与S3相连接，再经过3*3卷积，得到特征图记为Q3；Q3经过转置卷积上采样后与S2相连接，再经过3*3的卷积，得到的特征图记为Q2；

PAN自底向上传递强定位信息，特征图Q2作为底层特征R2，R2经过下采样后，与Q3相连接，得到特征图记为R3；R3经过下采样后与Q4相连接，得到的特征图记为R4；R2、R3、R4分别经过3*3的卷积，得到特征图T2、T3、T4；

在颈部Neck结构的后3个C3模块之后分别集成一个AMECA注意力模块，分别将特征图T2、T3、T4作为原始输入特征图，分别经过全局平均池化模块与全局最大池化模块，并将得到的两个特征图进行相加，压缩空间信息，随后使用1*1卷积学习通道注意力信息，得到的通道注意力信息与原始输入特征图结合，最终得到具体的通道注意力特征图D1、D2、D3；

输出部分：将特征图D1、D2、D3分别输入YOLOv5-MSTA检测头网络，最终得到检测结果。

如图2所示，所述的Mosaic9数据增强包括：首先从总数据集中取出一个批量的数据，每次从中随机取出9张图片，进行随机位置的裁剪、缩放，合成新图片；上述过程重复batch-size次，最后得到包括batch-size个经过了Mosaic9数据增强后图片的一个批量的新数据，再传递给神经网络进行训练。本发明采用Mosaic9数据增强，采用9张图像进行随机裁剪、缩放后，再随机排列拼接形成一张图片，这样改进实现丰富数据集的同时，增加了小样本目标，提升网络的训练速度和泛化能力，在进行归一化操作时会一次性计算9张图片的数据，因此模型对内存的需求降低。

具体的，所述的CBRM操作包括Conv、BN、ReLU和MaxPool。

具体的，所述的第一下采样层、第二下采样层、第三下采样层中均包括Shuffle_Block(d)模块，所述的Shuffle_Block(d)不再进行通过分流操作，将输入特征分为两个分支，左侧分支有2个卷积层，分别为步长为2的3×3深度卷积和1×1的普通卷积；右侧分支上有三个卷积层，分别为1×1的普通卷积、步长为2的3×3深度卷积和1×1的普通卷积；左右两侧分支通过Concat进行拼接来使左右两侧的特征进行融合，最后通过通道混洗操作来启用两个分支之间的信息通信。与Shuffle_Block(c)不同的是在左右两侧都引入了3×3的深度卷积来实现下采样。

具体的，所述的第二卷积归一化层、第三卷积归一化层、第四卷积归一化层中均包括Shuffle_Block(c)模块，所述的Shuffle_Block(c)模块对每个通道进行分流，划分为两个分支，根据要减少模型的碎片化程度准则，在左边的分支上没有进行任何操作，在右边的分支上有3个卷积层，分别是1×1普通卷积、3×3深度卷积、1×1普通卷积，1×1普通卷积、3×3深度卷积、1×1普通卷积三个卷积层拥有相同的输入和输出通道，其中两个所述的1×1Conv不再是组卷积而改变为普通的卷积，在所述的3个卷积之后，两个分支是通过Concat进行拼接；这样能够满足输入和输出通道是一样的，两个分支是通过Concat进行拼接的结果进行通道混洗操作，通过通道混洗操作启用两个分支之间的信息通信。

具体的，所述的转置卷积上采样的运算步骤包括：

(1)在输入特征图元素间填充s-1行、列0，其中s表示转置卷积的步距；

(2)在输入特征图四周填充k-p-1行、列0，其中k表示转置卷积的kernel_size大小，p为转置卷积的填充；

(3)将卷积核参数上下、左右翻转；

(4)做正常卷积运算，填充0，步距1。

采用转置卷积的上采样方式，通过网络自行学习，来获取最优的上采样方式，实现了语义级别的上采样，使特征含有更强的语义信息。转置卷积计算过程是将输入的每个元素值作为卷积核的权重，相乘后作为该元素对应的上采样输出，不同输入的重叠的输出部分直接相加作为输出。

所提出的AMECA注意力模块如图4所示，将特征图经过全局平均池化模块与全局最大池化模块，并将得到的两个特征图进行相加，压缩空间信息，随后使用1*1卷积学习通道注意力信息，得到的通道注意力信息与原始输入特征图结合，最终得到具体的通道注意力特征图。AMECA避免了降维，有效捕获了跨通道交互的信息，使得网络可以更精确地定位并识别到目标区域。

所述的AMECA注意力模块的流程包括：

(1)首先输入特征图X，特征图X的维度是H*W*C；

(2)对输入特征图X进行空间特征压缩；在空间维度，使用全局平均池化GAP，得到1*1*C的特征图F1；使用全局最大值池化GMP，得到1*1*C的特征图F2；

(3)将F1与F2融合得到1*1*C的特征图F3，获得更高层次的语义信息；

(4)对融合后的特征图F3进行通道特征学习；通过1*1*1卷积，学习不同通道之间的重要性，此时输出的特征图F4维度还是1*1*C；

(5)将特征图F4经过σ函数得到F41；

(6)最后是通道注意力结合，将通道注意力的特征图F41与原始输入特征图X，进行逐通道乘积，最终输出具有通道注意力的特征图X'；

其中，H，W和C分别表示输入特征图的高、宽和通道数，σ代表激活函数；

所述的特征图T2、T3、T4作为输入特征图X，分别得到的输出特征图X'为特征图D1、D2、D3。

另一种实施例，一种计算机可读介质其上存储有计算机程序用于执行如上所述的方法。

轻量化处理后的网络计算量和参数量大幅下降，在YOLOv5网络结构模型的Neck结构的最后3个C3模块之后引入AMECA注意力机制，调整模型空间和通道的信息提取方式。该方法很好地保证精度没有损失，可以满足齿轮表面缺陷实时检测的需求。

本发明做了对比试验，如表2所示。对比试验从参数量、计算量、模型大小进行了衡量，参数量、计算量和模型大小越小，说明网络复杂度越低。采用ShuffleNetv2模块的检测网络其参数量、计算量和模型大小均低于其他相同参数条件下的未采用ShuffleNetv2模块的检测网络，从而说明ShuffleNetv2模块具有轻量化的作用。

表2不同模型的参数量、计算量、模型大小对比

根据检测结果以及表中的结果得出，MSTA-YOLOv5模型相较于YOLOv3、YOLOv4和YOLOv5s模型具有很大的优势。相较于原始的YOLOv5s模型，模型、参数量和计算量都得到了大幅度减少，参数量约减少了46％，计算量减少了50％，模型大小约减少了44％。新的模型更加精简，复杂度明显降低，实现了在移动端部署的要求，对齿轮表面缺陷具有较好的检测结果。

上述虽然结合附图对发明的具体实施方式进行了描述，但并非对本发明保护范围的限制，在本发明的技术方案的基础上，本领域技术人员不需要付出创造性劳动即可做出的各种修改或变形仍在本发明的保护范围以内。