改进YOLOv5s的钢材表面缺陷检测算法

2024-04-26 21:03吕秀丽卢海滨侯春光王志刚
化工自动化及仪表 2024年2期
关键词:缺陷检测注意力机制深度学习

吕秀丽 卢海滨 侯春光 王志刚

基金项目:黑龙江省教育科学规划课题(批准号:GJB1421131)资助;黑龙江省高等教育教学改革研究项目(SJGY20210110)资助的课题。

作者简介:吕秀丽(1971-),教授,从事图像处理、模式识别领域的教学和科研工作。

通讯作者:卢海滨(1997-),硕士研究生,从事图像处理、模式识别的研究,18746933517@163.com。

引用本文:吕秀丽,卢海滨,侯春光,等.改进YOLOv5s的钢材表面缺陷检测算法[J].化工自动化及仪表,2024,51(2):

301-309.

DOI:10.20030/j.cnki.1000-3932.202402020

摘 要 为提高钢材表面缺陷检测的准确率,提出一种改进YOLOv5s的钢材表面缺陷检测算法。首先,在特征提取网络中引入Swin Transformer结构,增强网络对特征的感知能力;其次,添加坐标注意力机制,加强对重要特征信息的关注;最后,针对钢材缺陷的特点增加检测层,提升多尺度目标检测能力,并使用SIOU损失函数评估检测效果。将所提出的算法在公开数据集NEU-DET上进行消融实验,结果表明:所提算法能有效提高钢材表面缺陷目标检测的准确率。

关键词 缺陷检测 深度学习 改进YOLOv5s Swin Transformer 注意力机制

中图分类号 TP391.41  文献标志码 A   文章编号 1000-3932(2024)02-0301-09

钢材作为重要的工业材料之一,在船舶、航天及建造等诸多领域中有着不可替代的地位。在生产过程中难免会在其表面产生一些缺陷,如凹面、裂纹及划痕等[1]。这些缺陷有可能降低钢材的使用性能,進而导致事故发生的风险增加。因此,在钢材投入实际使用前,使用准确高效的检测方法对其表面进行缺陷检测确认其质量具有重要意义[2]。

机器视觉技术具有安全可靠、适用范围广及检测效率高等优点。近年来,国内外研究者在机器视觉技术上针对缺陷检测进行了广泛的研究,主要分为基于机器学习的方法和基于深度学习的方法。基于机器学习的方法主要依赖人工提取特征,难以捕捉高级语义特征,适用于简单场景,对于特征复杂、缺陷不规则的情况并不适用。与机器学习算法相比,基于深度学习的算法能够自动学习模式特征,并将特征学习融入到建立模型的过程中,从而削弱了人为设计特征导致的不完备性,在复杂环境中仍具有较强的鲁棒性。凭借着上述优点,深度学习技术迅速发展,在一些领域逐渐取代了传统机器学习方法。王立中和管声启针对传统钢材缺陷检测方法无法自动准确提取特征的问题,以卷积神经网络(CNN)为框架,建立了钢材表面缺陷检测的基础模型[3],通过fine-tuning的方法训练,实验证明该方法具有很强的抗噪性,但是该模型使用的图像数据量较少。唐寿权在进行样本扩充后,使用ResNet-50结构提取特征和R-FCN结构增强特征感知能力的方法改进了原始Faster R-CNN模型[4],有效提高了检测精度,但对于钢材缺陷中的小目标仍存在漏检的情况。李维刚等以YOLOv3算法为基础,使用K-means算法聚类出更适合的先验框参数,并增加了残差结构[5],提高了对钢材表面缺陷中小目标的检测能力,但改进后的模型参数量大幅度增加。袁洪强等利用知识蒸馏机制结合ResNet50和MobileNetV3,生成了一种轻量化的模型[6],模拟实际检测场景通过随机翻转与旋转的方式增强数据,模型训练后成功部署在移动端并具有较高的检测精度。马燕婷等将Transformer机制融入YOLOv5中[7],提升了主干网络特征提取的全局性,并改进BiFPN网络为T-BiFPN网络结构,使得图像的特征融合能力得到了增强,使用RepVGG网络替换部分卷积层,提高了特征提取能力,有效提升了钢材表面缺陷的检测能力。

笔者基于YOLO[8]系列经典目标检测算法中的YOLOv5s,提出了一种改进YOLOv5s的钢材表面缺陷检测算法(YOLOv5s-STCD),在公开数据集NEU-DET[9]上证明了该算法对于钢材表面缺陷有较好的检测能力。

1 算法原理

YOLOv5是目前广泛应用的单阶段目标检测算法之一,具有检测精度高、速度快的特点,按照网络深度和宽度的加深程度可分为YOLOv5s、YOLOv5m、YOLOv5l和YOLOv5x 4种模型。笔者选用模型体积较小的YOLOv5s作为基准模型,YOLOv5s网络结构如图1所示。

YOLOv5s网络分为四大模块:Input(输入)、Backbone(主干网络)、Neck(特征融合网络)和Head(预测头部)。Backbone部分的作用是对图像特征进行提取,由CBS、CSP1_X和SPPF组成(其中,SPPF结构在YOLOv5-6.0中首次使用,作为SPP结构的优化形式,同样能够实现输出自适应尺寸的功能,但SPPF改变了结构形式,减少了模型计算量,提升了模型速度)。Neck部分将Backbone提取的特征进行融合,由CSP2模块、上采样模块和Concat(联合模块)构成。Head部分包含3个不同尺度的特征检测器,对输出的特征图进行目标预测。

2 YOLO-STRCD算法介绍

笔者提出的YOLOv5s-STCD的整体结构框架如图2所示。算法YOLOv5s-STCD中,ST代表Swin Transformer[10]网络结构,C代表坐标注意力机制Coordinate Attention[11],D代表添加的检测层Detection。在主干网络中使用C3STR结构替换部分C3结构(CSP+CBS),融入滑动窗口的多头自注意模块,增强网络对特征的感知能力;为减少无用信息的干扰,引入坐标注意力机制,提高感兴趣信息关注度;针对数据集中大部分缺陷目标尺度大小不均衡这一特性,增加一个检测层,并在Neck部分添加相应的特征融合结构;改进边界框回归损失函数,采用SIOU损失函数,加快推理速度并进一步提升模型检测钢材表面缺陷的精度。

2.1 C3STR结构

Transformer凭借自注意力机制以及良好的模态融合能力,在自然语言处理(Natural Language Processing,NLP)任务中取得突破性进展,但其从NLP领域跨越到计算机视觉(Computer Vision,CV)领域存在一定困难,例如文本单词与高分辨率图像像素之间有巨大的差异。Swin Transformer通过构建层次化特征图和滑动窗口的形式解决了这些差异,在目标分类和目标检测任务上都取得了极佳的效果。为了增强特征提取能力,在主干网络部分将Swin Transformer中的窗口自注意力融入C3模块形成C3STR结构,如图3所示。

Swin Transformer Block由窗口多头自注意层(Window Multi-head Self-attention,W-MSA)、移位窗口多头自注意层(Shifted-window Multi-head Self-attention,SW-MSA)、多层感知器(Multilayer Perceptron,MLP)组成,各模块间连接方式采用残差结构,W-MSA与SW-MSA成对出现[12],多头自注意力机制计算方式如下:

Attention(Q,K,V)=SoftMAX(QK/+B)V (1)

式中 B——相对位置偏置;

d——Query/Key的维度;

K——Key(索引)矩阵;

Q——Query(查询)矩阵;

V——Value(值)矩阵。

相比于传统的Transformer,C3STR采用滑动窗口的形式,滑动窗口使其包含不同的像素点以实现相邻窗口间的信息传递,减小了计算量,可提高网络的运算效率。

2.2 引入坐标注意力机制

注意力机制凭借简单有效和计算量增幅小的特点,在深度学习领域中引起许多学者的兴趣并展开了研究。例如压缩和激活(Squeeze and Excitation,SE)[13]注意力机制,使用全局平均池化操作提取通道信息,通过两个全连接层构建通道间的联系并得到各通道的权重值,但其忽略了空间特征;卷积块的注意力模块(Convolutional Block Attention Module,CBAM)[14]通过引入大尺度的卷积核提取空间特征,增加了对空间信息的关注,但其仅仅考虑了局部区域的信息,而无法建立远距离的依赖;坐标注意力机制(Coordinate Attention,CA)具有捕获跨通道信息的能力,同时能够加强对位置和方向信息的感知能力,并且其能够简单地插入到轻量级网络的核心模块中。

在钢材缺陷图像中,裂纹、凹面等类型缺陷容易受到背景因素的影响,笔者在模型中引入CA来抑制此种情况,原理如图4所示。为了避免空间信息全部压缩到通道中,CA分别在宽度和高度两个方向上进行全局平均池化,可以得到高度与宽度方向上的两组特征图,将二者拼接后通过激活函数获取两个维度上的权值。CA模块主要分为坐标信息嵌入和坐标注意力生成两部分。

图4 CA模块

C——特征图通道数;  H——输入特征图的高度;

W——输入特征图的宽度; r——通道压缩率

坐标信息嵌入。由于全局平均池化将全局空间信息压缩到通道描述符中,忽略了位置信息,故将其转换为一对一维特征编码操作,高度为h的第c通道的输出可以表示为:

z(h)=x(h,i)    (2)

式中 x(h,i)——水平方向第c通道的输入。

同样,宽度为w的第c通道的输出可以写成:

z(w)=x(j,w)     (3)

式中 x(j,w)——垂直方向第c通道的输入。

上述两种变化分别沿两个空间方向聚合特征,得到一对方向感知的特征图。

坐标注意力生成。通过信息嵌入变换后,将上面的变换进行连接操作,然后使用1×1卷积变换函数F1对其进行变换操作:

f=δ(F[z,z])      (4)

式中 δ——Sigmoid激活函数。

将f沿着空间维数分解为2个单独的张量,将其变换后展开得到CA的输出y:

y(i,j)=x(i,j)×g(i)×g(j)     (5)

式中 g——水平方向上的注意力权重;

g——垂直方向上的注意力权重。

由此,g、g两个权重被添加到输入特征图中来增强特征图的表示能力。

为更好地验证不同注意力机制对YOLOv5s网络模型的影响,将添加了3种注意力机制的YOLOv5s模型在NEU-DET数据集上进行对比实验,结果见表1。实验采用召回率R、精确度P和均值平均精度mAP进行评价。实验表明添加3种注意力机制对模型性能均有不同程度的提升,CA表现最佳,CBAM次之,SE最差。

表1 对比实验     %

2.3 增加检测层

原始YOLOv5s模型包含3个不同尺度的检测层,当输入图像大小为640×640时,检测层输出的3个尺度为80×80、40×40、20×20,根据感受野大小不同进行区分,大尺寸特征图检测小目标,小尺寸特征图检测大目标。在对钢材缺陷进行检测的实际应用场景中,各缺陷特征的尺寸均不相同,所以映射到图像中的钢材表面缺陷特征有許多大小不同的尺度。为了提高YOLOv5s模型对不同尺寸钢材缺陷目标的检测能力,在原始模型基础上增加一个检测层。在主干网络中最后一个C3层后添加一组CBS和C3模块,同时在特征融合部分对应位置上添加一组上采样和相应的下采样,这样能更好地实现特征融合功能。新增的检测层有利于增强网络的泛化能力。

2.4 改进边框回归损失函数

目标检测领域中应用较广的损失函数有IOU(Intersection over Union) Loss[15]、GIOU(Generalized Intersection over Union) Loss[16]、DIOU(Distance Intersection over Union) Loss[17]、CIOU(Complete Intersection over Union)[17]等。IOU是较早提出来用预测框和真实框的重合度来判断模型预测准确度的指标;GIOU在IOU的基础上解决了两框无重叠导致的梯度为零的问题。DIOU在GIOU的基础上将两矩形框中心点距离加入,改进后训练更加稳定,收敛更快。CIOU将两矩形框宽高比加入计算:如果两框的重叠区域较小,则宽高比在损失函数中影响较小,此时着重优化两框间的距离;反之,如果重叠区域较大,此时着重优化两框的宽高比。但以上几种损失函数均未考虑两框之间的方向是否匹配,导致收敛速度缓慢且效率较低。鉴于上述问题,GEVORGYAN Z提出了损失函数SIOU,其考虑了回归向量的角度,增加了角度惩罚项,提高了模型的收敛速度和推理的准确性。

SIOU由4个代价函数构成,分别是角度代价、距离代价、形状代价和IOU代价。添加角度感知LF组件,有效减少了与距离相关的变量数量,添加角度感知LF组件的计算方式如下:

Λ=1-2×sin2arcsin α-     (6)

式中 α——预测框和真实框中心点连线的水平夹角;

Λ——角度损失。

根据上面的角度代价重新定义了距离代价函数:

Δ=(1-e)       (7)

γ=2-Λ

当t为x时,ρ表示真实框和预测框宽的差值比真实框和预测框最小外接矩形宽的平方值;当t为y时,ρ表示真实框和预测框高的差值比真实框和预测框最小外接矩形高的平方值。

形状代价函数定义如下:

Ω=(1-e)     (8)

ω=      (9)

ω=

式中 w、h——预测框的宽、高;

w、h——真实框的宽、高;

θ——用于控制对形状损失的关注程度,为了避免过于关注形状损失而降低对预测框的移动,θ的取值范围应在[2,6]。

IOU代价函数仍按照两框之间的重合度计算。综上,SIOU损失函数表达式为:

LOSS=1-IOU+     (10)

3 实验结果与分析

3.1 数据集介绍

在公开数据集NEU-DET上进行钢材表面缺陷检测实验,NEU-DET数据集包含裂纹(Crazing,Cr)、夹杂物(Inclusion,In)、斑块(Patches,Pa)、凹面(Pitted Surface,PS)、轧制氧化皮(Rolled-in Scale,RS)和划痕(Scratches,Sc)共6种缺陷。NEU-DET数据集的部分缺陷图片如图5所示。数据集共有1 800张圖片以及相对应的XML格式注释文件。将XML格式的注释文件转换为YOLO模型能够识别的TXT格式,将其与图片数据对应整合之后,按照6∶2∶2的比例对6个类别的缺陷数据进行均匀随机划分,得到训练集图片1 080张,验证集图片360张,测试集图片360张。

3.2 实验设备及参数设置

实验在Windows11操作系统下进行,内存为64 GB,深度学习框架为PyTorch1.12.1,Python版本为3.8.13,CUDA版本为11.6,CPU为Intel(R)Core(TM) i7-12700H@2.30 GHz,GPU为NVIDIA

GeForce RTX 3060 6GB。在训练中,超参数设置为初始学习率0.001,衰减系数为0.000 5;动量为0.937,批次为16,训练次数为300。

3.3 实验评估指标

使用多个指标来对算法进行评估,主要通过平均精度AP和均值平均精度mAP[18]分别对各类缺陷以及整体网络模型的性能进行评价,同时需要考虑精确度P、召回率R及模型体积等参数。

3.4 实验结果分析

为检验笔者提出算法的有效性,在公开数据集NEU-DET上进行了消融实验,实验结果见表2,其中Detection表示检测层。

表2 消融实验结果 %

由表2可知,以上改进方法中,在网络中加入CA,mAP提升了0.7%;增加检测层,mAP提升了0.5%;引入C3STR结构,mAP提升了0.8%。当原YOLOv5s中加入全部改进机制时,整体mAP提升了5.0%。

将YOLOv5s模型的训练损失与YOLOv 5s-STCD的训练损失曲线进行对比,如图6所示。

图6 改进前后模型损失曲线对比

由图6可知,改进后的模型YOLOv5s-STCD与原始YOLOv5s相比,在定位损失、置信度损失和分类损失3个方面均有所改善,在前15次训练中,两种模型损失值接近,损失曲线近似重合;此后差异开始逐渐明显,YOLOv5s-STCD模型的损失开始加速收敛,二者曲线开始分离,并且改进后的模型损失值始终低于原始YOLOv5s模型;在260个训练次数后,二者损失曲线均趋向于平稳,在训练结束时YOLOv5s-STCD模型的3种损失均低于YOLOv5s模型,证明改进方案对YOLOv5s模型具有优化作用。

为更直观地展示改进后算法的检测效果,将部分检测结果与原YOLOv5s模型的检测结果进行对比,如图7所示。

由图7所示的改进前后模型检测结果的对比可知,改进后的模型改善了原始模型检测时存在漏检的情况,并且提高了原始模型对部分缺陷目标检测时的置信度,在PS类型缺陷上表现最为明显,有效地证明了改进模型相对于原始模型的性能提升。

3.5 对比实验

为了进一步证明笔者算法在钢材表面缺陷检测方面的优势,选取了4种主流目标检测算法SSD300[19]、Faster-RCNN[20]、YOLOv4[21]和

YOLOv5s在NEU-DET数据集上与笔者算法对比,结果见表3。

由表3可知,相比于SSD300、YOLOv4、YOLOv5s目标检测算法,Faster-RCNN凭借着两阶段检测网络的优势在钢材缺陷检测中精度最高,但其缺点也很明显,模型体积较大。从检测精度上看,SSD300在大目标类型Ps缺陷中效果较好,但在以小目标为主的Sc类型上效果较差,这是由于SSD300算法对于小目标的检测存在些许不足。YOLOv4在上述实验中的mAP值为69.3%,存在检测精度低、模型体积大的问题。YOLOv5s-STCD作为改进后的模型,比原YOLOv5s的mAP值提高了5.0%。综上所述,笔者所提出的改进算法在精度和模型体积上更具有优势。

4 结束语

针对钢铁表面缺陷检测问题,提出一种YOLOv5s-STCD算法。引入Swin Transformer结构增强特征提取能力;添加CA注意力机制,增加重要特征的关注度;增加目标检测层,提高对部分类型目标的检测精度;改进损失函数。在实验中,笔者所提算法相较于原YOLOv5s,均值平均精度mAP在NEU-DET数据集上提升了5.0%;与主流算法SSD300相比,mAP值提升了6.4%,表明笔者所提算法对于钢材表面缺陷检测的有效性。今后将进行模型的轻量化研究,继续优化模型,提高检测效率,尝试将其布署到移动端。

参 考 文 献

[1] 李鑫,汪诚,李彬,等.改进YOLOv5的钢材表面缺陷检测算法[J].空军工程大学学报(自然科学版),2022,23(2):26-33.

[2] 朱宏平.基于卷积神经网络的钢表面缺陷检测方法[J].工业控制计算机,2020,33(8):83-84.

[3] 王立中,管声启.基于深度学习算法的带钢表面缺陷识别[J].西安工程大学学报,2017,31(5):669-674.

[4] 唐寿权.基于深度学习的带钢产品表面缺陷检测方法[D].沈阳:东北大学,2018.

[5] 李维刚,叶欣,赵云涛,等.基于改进YOLOv3算法的带钢表面缺陷检测[J].电子学报,2020,48(7):1284-1292.

[6] 袁洪强,杜国锋,余泽禹,等.基于轻量神经网络的钢材表面缺陷快速识别[J].科学技术与工程,2021,21(34):14651-14656.

[7] 马燕婷,赵红东,阎超,等.改进YOLOv5网络的带钢表面缺陷检测方法[J].电子测量与仪器学报,2022,36(8):150-157.

[8] REDMON J,DIVVALA S,GIRSHICK R,et al.You only look once: Unified, real-time object detection[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition(CVPR).Piscataway,NJ:IEEE,2016.DOI:10.1109/CVPR.2016.91.

[9] HE Y,SONG K,MENG Q,et al.An end-to-end steel surface defect detection approach via fusing multiple hierarchical features[J].IEEE Transactions on Instrumentation and Measurement,2019,69(4):1493-1504.

[10] LIU Z,LIN Y Y,CAO Y,et al.Swin transformer:Hierarchical vision transformer using shifted windows[C]// Proceedings of the IEEE/CVF International Conference on Computer Vision.Piscataway,NJ:IEEE,2021.DOI:10.1109/ICCV48922.2021.00986.

[11] LI B Y,LIU Y,WANG X G.Gradient harmonized sin-gle-stage detector[C]//Proceedings of the AAAI Conference on Artificial Intelligence.Honolulu:AAAI,2019.DOI:10.48550/arXiv.1811.05181.

[12] 趙文清,康怿瑾,赵振兵,等.改进YOLOv5s的遥感图像目标检测[J].智能系统学报,2022,18(1):86-95.

[13] 胡欣,周运强,肖剑,等.基于改进YOLOv5的螺纹钢表面缺陷检测[J].图学学报,2023,44(3):427-437.

[14] 江培营,陶青川,艾梦琴.基于注意力机制和深度学习的钢板表面缺陷图像分类[J].计算机应用与软件,2021,38(9):214-219.

[15] 唐靚,余明慧,武明虎,等.基于改进 YOLOv5 的绝缘子缺陷检测算法[J].华中师范大学学报(自然科学版),2022,56(5):771-780.

[16] REZATOFIGHI H,TSOI N,GWAK J Y,et al.Gener alized intersection over union:A metric and a loss for bounding box regression[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition.Piscataway,NJ:IEEE,2019.DOI:10.1109/ CVPR.2019.00075.

[17] ZHENG Z,WANG P,LIU W,et al.Distance-IoU loss:Faster and better learning for bounding box regression[C]//Proceedings of the AAAI Conference on Artificial Intelligence.NewYork:AAAI,2020.DOI:10.1609/aaai.v34i07.6999.

[18] 赵月,张运楚,孙绍涵,等.基于深度学习的螺纹钢表面缺陷检测[J].计算机系统应用,2021,30(7):87-94.

[19] LIU W,ANGUELOV D,ERHAN D,et al.Ssd:Single shot multibox detector[C]//Computer Vision-ECCV 2016:14th European Conference.Amsterdam,2016.

[20] REN S Q,HE K M,GIRSHICK R B,et al.Faster R-CNN:Towards real-time object detection with region proposal networks[J].IEEE Transactions on Pattern Analysis and Machine Intelligence,2017,39(6):1137-1149.

[21] BOCHKOVSKIY A,WANG C Y,LIAO H Y M.YOLOv4 :Optimal speed and accuracy of object detection[J].arXivPreprint,2020.DOI:10.48550/arXiv.2004.1 0934.

(收稿日期:2023-03-21,修回日期:2024-01-30)

Improved YOLOv5 Algorithm for Steel

Surface Defect Detection

LV Xiu-li1, LU Hai-bin1, HOU Chun-guang2, WANG Zhi-gang2

(1. School of  Physics and Electronic Engineering, Northeast  Petroleum University;

2. Office of  the Network Security and  Information  Technology  Committee of

the CPC  Daqing Municipal  Party Committee)

Abstract   For purpose of improving the accuracy in detecting steel surface defects,an improved YOLOv5s algorithm for the steel surface defect detection was proposed. Firstly, having Swin Transformer structure introduced into the feature extraction network to enhance feature perception ability of the network; secondly, having coordinate attention mechanism employed to enhance the attention to important feature information;finally, through considering the characteristics of steel defects, having detection layer added to improve the  detection ability of multi-scale target, and having SIOU lossfunction used to evaluate the detection effect.Im plementing ablation experiments of the proposed algorithm on the public data set NEU-DET shows that,the proposed algorithm can effectively improve the average accuracy in detecting steel surfaces target defect.

Key words   defect detection, deep learning, improved YOLOv5s, Swin Transformer, attention mechanism

猜你喜欢
缺陷检测注意力机制深度学习
基于深度学习的问题回答技术研究
基于LSTM?Attention神经网络的文本特征提取方法
基于注意力机制的双向LSTM模型在中文商品评论情感分类中的研究
InsunKBQA:一个基于知识库的问答系统
超声波与特高频方法下的GIS局部放电检测技术分析
基于极限学习机的玻璃瓶口缺陷检测方法研究
MOOC与翻转课堂融合的深度学习场域建构
大数据技术在反恐怖主义中的应用展望
深度学习算法应用于岩石图像处理的可行性研究
基于深度卷积网络的人脸年龄分析算法与实现