基于YOLOv5 改进的铝片表面缺陷检测
2024-01-16谢延楠刘兴德
谢延楠,刘兴德
(1.吉林化工学院 信息与控制工程学院,吉林吉林,132022;2.吉林化工学院 机电工程学院,吉林吉林,132022)
0 引言
全球各方面的发展不断加快,铝制品的需求量不断增加,铝片的质量对于铝制品的性能具有重要影响。因此,对铝片表面缺陷进行有效检测具有重要的实际意义。传统的缺陷检测有准确性差和检测速度慢的问题。随着计算机视觉技术的快速发展,使用深度学习的目标检测方法成为主流。
随着其精度和速度不断提升,服务于目标检测需求的神经网络模型也越来越多,众多的网络模型为工厂车间内的工件检测提供了新的可能。常见的目标检测算法分为两大类,单阶段检测算法和双阶段检测算法[1],单阶段算法以SSD、YOLO 系列为典型,而双阶段则以Faster R-CNN 网络为主。其中基于 YOLOv5 的检测方法在速度和准确性上具有一定优势。
针对表面缺陷检测问题,王子豫等[2]人提出基于改进的YOLOv5 算法的零件表面缺陷检测技术,引入CBAM 注意力机制,并融合BiFPN 网络与可变性卷积策略,并采取剪枝方法提升模型性能。杜少聪等[3]人提出基于改进YOLOv5 的钢轨表面缺陷检测技术,将多头注意力层嵌入到骨干网络末端,并引入全局依赖关系,构建跨层加权级联结构,将浅层网络融入深层网络中,提高检测性能。吴健生[4]等人提出基于改进Faster R-CNN 算法的带钢缺陷检测,在主干网络采用可变形卷积模块,FPN 多尺度检测模块以及在RPN 网络中融合CBAM 模块,提高检测精度。刘艳菊等[5]人提出基于改进SSD 的工件表面缺陷检测,采用了以ResNet 替换SSD 中原始的VGGNet 的方法,研究了小目标检测的问题;采用了对深层特征进行反卷积且将深层特征与浅层特征融合的方法,提升网络整体性能。
针对工业铝片缺陷数据集,本文研究一种基于YOLOv5的网络模型,嵌入不同的卷积注意力机制模块以增强图像的特征提取能力,之后修改YOLOv5 本身的CIOU 损失函数,降低预测框的回归损失。
1 YOLOv5 算法概述
YOLOv5 是一种高效的目标检测算法,由YOLOv4 算法改进而来。YOLOv5 算法具有检测速度快、实时性高、鲁棒性好等优点,广泛应用于工业检测领域。YOLOv5 作为YOLO 系列最常用的模型,提供了四种大小的模型。其中YOLOv5s 的模型结构最小,其他几种模型结构相比更为复杂,但平均精度会有所提升。
图1 为YOLOv5-6.0 版本的模型结构。
图1 网络结构图
下面分别对四个部分进行分解说明:
(1)Input 输入端:YOLOv5 在输入阶段结合Mosaic 数据增强和自适应图像缩放,增强推理能力,提高网络效率,使其能够很好地适应各种输入。
(2)Backbone 主干网络:相比于Yolov5-5.0 版本,Yolov5-6.0 版本使用卷积核代替Focus 模块,便于模型的导出,且效率更高。该部分三个部分组成,分别是CBS、CSP1_X、SPPF,主要负责对输入图像进行特征提取。三个模块的详细信息如图1 中所示。
(3)Neck 颈部网络:Neck 部分主要是对之前提取的特征进行整理使其能够被更好地利用从而提取出更有效的特征,采用基于多特征图融合的FPN+PANet 结构,从自顶向下和自底向上两个方向出发,同时增强了语义信息和位置信息的提取,同时设计了的CSP2_X 结构,进一步加强了网络的特征融合能力。
(4)Head:采用CIOU Loss 作为模型的损失函数,生成边界框和类别预测,同时采用非极大值抑制NMS 的方法,筛选保留最佳的目标框,完成预测任务并记录缺陷位置及类别[6]。
2 改进YOLOv5 网络模型
■2.1 注意力机制
为提高检测精度,提升对定位内容的关注,加入注意力机制模块对模型的性能进行优化,本文使用的是通道注意力和空间注意力的集成者CBAM[7],这是一种用于前馈卷积神经网络的简单而有效的注意模块[8],CBAM 为轻量级通用模块,即插即用。
CBAM 由通道和空间注意力两个子模块组成,图2 为CBAM 模块的结构图。
图2 CBAM 模型结构图
在通道注意力模块(CAM)的结构中,这个过程的目标是通过对输入特征图F 进行不同池化方式的处理,然后通过共享的神经网络来捕捉全局信息,并通过逐元素相加和激活函数来生成通道注意力特征图Mc。这个特征图可以帮助网络更好地聚焦于输入数据中最重要的通道信息,从而改善模型性能。CAM 的模型结构图如图3 所示。
图3 CAM 模型结构图
空间注意力模块(SAM),在SAM 中,输入是通道注意力模块得到的特征图Mc 和输入特征图 F 进行乘法处理后的结果F'。下面是SAM 的步骤:
首先特征图 F' 分别进行MaxPooling 和AvgPooling 操作,产生两个池化特征图。然后,将处理得到的池化特征图进行通道拼接,将它们在通道维度上连接起来后,将合并后的特征图输入一个卷积操作,将通道数量减少到1。对卷积输出的特征图进行 sigmoid 激活,生成范围在0 到1 之间的空间注意力特征图 Ms[9],图4 为SAM 的结构图。
图4 SAM 模块结构图
将得到的空间注意力特征图Ms 与SAM 的输入特征图F'进行逐元素乘法操作,最终得到生成的特征。这样的操作将融合通道和空间注意力的信息。
总之,通道注意力模块得到的特征图F'执行池化、通道拼接、卷积和激活操作,生成一个用于强调输入特征图中重要空间位置的SAM 特征图 Ms。最终,通过将Ms 与该模块输入特征图F'相乘,得到一个融合了通道和空间的注意力信息的最终特征。这有助于网络更好地关注输入数据中的重要通道和空间信息,从而提升模型性能。
■2.2 损失函数
损失函数的作用就是判断预测缺陷和真实缺陷之间的差距,若预测缺陷信息越接近真实缺陷信息,则损失函数值越小。YOLOv5 的损失主要包含三个方面:边界框损失(bbox_loss)、分类损失(cls_loss)、置信度损失(obj_loss)。总损失的表达式为三个损失的和:
边界框损失,关注的重点是预测边界框和真实边界框的交集比例,用交并比即IOU 表示[10],其中设预测边界框为A,真实边界框为B,则IOU 表示为:
当预测边界框与真实边界框完全重叠时,也就是说它们具有相交部分,这时候我们可以说它们的IOU(交并比)为1。相反地,如果预测边界框与真实边界框没有相交部分,那么它们的IOU 为0。在YOLOv5 网络中,默认使用CIOU Loss 来计算边界框的损失。相比于没有改进的IOU 损失,CIOU 损失考虑了边界框的尺寸信息,并综合考虑了边界框之间的位置关系。具体而言,CIOU 损失通过计算预测边界框与真实边界框之间的距离,将这一距离作为损失函数的一部分,用于优化目标检测算法[11],以下为CIOU 的损失计算公式:
其中ρ为两个边界框中心点间的欧氏距离,b和bgt为预测边界框的中心点和真实框的中心点,c 为预测边界框和真实框的最小闭包区域的对角线距离,α为权重参数,v 是用来衡量宽高比的一致性[13]。
分类损失,原网络分类损失函数使用二元交叉熵函数,该函数的定义为:
其中y 为输入样本对应标签,p 为预测样本为正样本的概率。
置信度损失,一个预测边界框对应一个置信度值,该值表示接近真实边界框的程度,置信度值越高,与真实框越吻合。计算置信度损失的函数是二元交叉熵函数。
本文使用的是Alpha-IoU Loss,基于现有IoU Loss 的统一幂化,即对所有的IoU Loss,增加α幂,当α等于1时[13],则回归到原始各个Loss 中:
■2.3 改进后的模型
将CBAM 注意力机制模块添加至Neck 部分,使网络能够从通道和空间两个方面提升提取特征的能力,并且替换Alpha-IoU Loss 的损失函数,提高性能,更适用于本数据集检测,图5 为更改后的网络结构模型。
图5 改进后的网络模型结构
图6 缺陷数据集
3 实验及结果分析
■3.1 实验数据
本文所使用的实验数据来自飞桨训练平台上的公开数据集,铝片表面工业缺陷检测数据集,共412 张韩有缺陷的铝片数据,共包含4 类缺陷(脏污dirty,褶皱fold,针孔pinhole,擦伤scratch),随机对该数据集进行划分,其中80%为训练集和20%为测试集,使用LabelImg 打标签软件对所得数据进行打标签处理。
本文中的训练使用的环境配置为:操作系统:Ubuntu 20.04.6 LTS;CPU:Intel(R) CPU @ 2.20GHz;GPU:NVIDIA Tesla T4(16G);Python:3.10;深度学习框架:Pytorch 2.0.1+cu118。
■3.2 性能分析
本文实验使用平均精度均值mAP 来评价模型的准确性和FPS 用来评估检测的速度,即每秒内可以处理的图片数量。
3.2.1 对比试验
对比试验可以通过和其他网络模型进行对比,更好地体现本文中改进模型的优势,本文实验将YOLOv5-αC(添加CBAM 注意力机制和Alpha-IOU 损失函数)与YOLOv3、Faster-RCNN、SSD、YOLOv7 以 及YOLOv5 模块添加不同种注意力机制模型进行对比,采用同样的训练集和测试集进行实验,不同模型对该数据集的性能对比如表1 所示。
表1 不同模型对比实验
YOLOv5 模型属于轻量级模型,并且简单易上手,检测速度和检测精度相对而言有所提高。虽然Faster-RCNN 网络模型的平均检测精度很高,但是检测速度很慢,无法做到检测的实时性,因此不考虑该网络模型。
3.2.2 消融实验
该实验内容的主要是明确地对比出有无注意力机制模块和替换损失函数对检测效果的影响,如表2 所示。
表2 消融实验结果
为了明确地看出模型改进的效果,采用测试集的部分图片进行测试,测试集图片在改进模型上的检测效果如图7所示。
图7 检测结果对比
4 结论
提高铝片工业缺陷的检测精度和检测速度在生产过程中具有重要意义,针对铝片表面工业缺陷问题,本文提出的一种添加CBAM 注意力机制和替换损失函数的改进YOLOv5网络能够有效地提高检测精度和速度。实验表明,与原模型比较,平均检测精度mAP 上涨3.1%,模型检测速度上涨27.784 帧/秒,符合工业实时检测要求。