改进YOLOv8的轻量化轴承缺陷检测算法

doi:10.3778/j.issn.1002-8331.2403-0444

摘要/Abstract

摘要： 针对轴承表面缺陷检测存在小目标检测精度低、模型复杂的问题，提出一种改进YOLOv8的轻量级轴承缺陷检测算法YOLO-SSW。在主干网络中添加3-D注意力机制SimAM，在不增加参数的前提下，使模型更加关注轴承表面缺陷特征的提取和表达；在颈部网络中嵌入C2f_SCConv模块，以减少空间维度和通道维度上的特征冗余，降低模型的计算负载；添加新的小目标检测层用于检测轴承表面的小尺寸缺陷，提高模型对小目标的检测能力；使用基于动态非单调聚焦机制的Wise-IoU作为边界框回归损失函数，加快网络的收敛速度。实验结果表明，改进算法的mAP达到了91.5%，比原算法提升了2.9个百分点，模型参数量仅为2.7×106，计算量为11.8 GFLOPs，在提高检测精度的同时节约了计算资源。

关键词: 缺陷检测, 轻量级网络, YOLOv8, SimAM

Abstract: Aiming at the problems of low accuracy of small targets, complex models, and difficulty in deploying edge devices for bearing surface defect detection, a lightweight bearing defect detection algorithm YOLO-SSW with improved YOLOv8 is proposed. Firstly, the 3-D attention mechanism SimAM is added to the backbone network to make the model pay more attention to the extraction and expression of bearing surface defect features without increasing parameters. Secondly, the C2f_SCConv module is embedded in the neck network to reduce feature redundancy in both spatial and channel dimensions and decrease the computational load of the model. Then, a new small target detection layer is added for detecting small-size defects on the bearing surface to improve the ability of the model to detect small targets. Finally, Wise-IoU based on a dynamic non-monotonic focusing mechanism is utilized as the bounding box regression loss function to accelerate the convergence speed of the network. The experimental results show that the mAP of the improved algorithm is 91.5%, which is 2.9% higher than the original algorithm. The model has only 2.7 million parameters and 11.8 GFLOPs of computation, which not only improves detection accuracy but also saves computational resources.

Key words: defect detection, lightweight network, YOLOv8, SimAM

姚景丽, 程光, 万飞, 朱德平. 改进YOLOv8的轻量化轴承缺陷检测算法[J]. 计算机工程与应用, 2024, 60(21): 205-214.

YAO Jingli, CHENG Guang, WAN Fei, ZHU Deping. Improved Lightweight Bearing Defect Detection Algorithm of YOLOv8[J]. Computer Engineering and Applications, 2024, 60(21): 205-214.

参考文献

[1] XU C, LI L, LI J, et al. Surface defects detection and identification of lithium battery pole piece based on multi-feature fusion and PSO-SVM[J]. IEEE Access, 2021, 9: 85232-85239.
[2] 黄广俊, 邓元龙. 融合改进LBP和SVM的偏光片外观缺陷检测与分类[J]. 计算机工程与应用, 2020, 56(22): 251-255.
HUANG G J, DENG Y L. Polarizer visual defect detection and classification based on improved LBP and SVM algorithm[J]. Computer Engineering and Applications, 2020, 56(22): 251-255.
[3] HE X, WU L, SONG F, et al. Research on fabric defect detection based on deep fusion DenseNet-SSD network[C]//Proceedings of the 2020 International Conference on Wireless Communication and Sensor Networks, 2020: 60-64.
[4] ZHANG D, ZHENG Z, LI M, et al. CSART: channel and spatial attention-guided residual learning for real-time object tracking[J]. Neurocomputing, 2021, 436: 260-272.
[5] HE J, WANG Y C, WANG Y T, et al. A lightweight road crack detection algorithm based on improved YOLOv7 model[J]. Signal, Image and Video Processing, 2024, 18: 847-860.
[6] ZHAO S, LI G, ZHOU M, et al. YOLO-CEA: a real-time industrial defect detection method based on contextual enhancement and attention[J]. Cluster Computing, 2023, 27: 2329-2344.
[7] YANG Q, MA S, GUO D, et al. A small object detection method for oil leakage defects in substations based on improved Faster-RCNN[J]. Sensors, 2023, 23(17): 7390.
[8] 张涛源, 谢新林, 谢刚, 等. 融合Transformer的带钢缺陷实时检测算法[J]. 计算机工程与应用, 2023, 59(16): 232-239.
ZHANG T Y, XIE X L, XIE G, et al. Real-time strip steel defect detection algorithm fused with Transformer[J]. Computer Engineering and Applications, 2023, 59(16): 232-239.
[9] ZHAO Y, CHEN B, LIU B, et al. GRP-YOLOv5: an improved bearing defect detection algorithm based on YOLOv5[J]. Sensors, 2023, 23(17): 7437.
[10] LI J, WEN Y, HE L. SCConv: spatial and channel reconstruction convolution for feature redundancy[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023: 6153-6162.
[11] YANG L, ZHANG R Y, LI L, et al. Simam: a simple, parameter-free attention module for convolutional neural networks[C]//Proceedings of the International Conference on Machine Learning, 2021: 11863-11874.
[12] TONG Z, CHEN Y, XU Z, et al. Wise-IoU: bounding box regression loss with dynamic focusing mechanism[J]. arXiv:2301.10051, 2023.
[13] LI Y, YAO T, PAN Y, et al. Contextual transformer networks for visual recognition[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022, 45(2): 1489-1500.
[14] LI X, WANG W, HU X, et al. Selective kernel networks[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2019: 510-519.
[15] SRINIVAS A, LIN T Y, PARMAR N, et al. Bottleneck transformers for visual recognition[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021: 16519-16529.
[16] REZATOFIGHI H, TSOI N, GWAK J Y, et al. Generalized 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, 2019: 658-666.
[17] ZHANG Y F, REN W, ZHANG Z, et al. Focal and efficient IOU loss for accurate bounding box regression[J]. Neurocomputing, 2022, 506: 146-157.
[18] 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, 2020: 12993-13000.
[19] LIU W, ANGUELOV D, ERHAN D, et al. SSD: single shot multibox detector[C]//Proceedings of the 14th European Conference on Computer Vision (ECCV 2016), Amsterdam, The Netherlands, October 11-14, 2016: 21-37.
[20] REN S, HE K, GIRSHICK R, et al. Faster R-CNN: towards real-time object detection with region proposal networks[C]//Advances in Neural Information Processing Systems, 2015.
[21] LI C, LI L, JIANG H, et al. YOLOv6: a single-stage object detection framework for industrial applications[J]. arXiv: 2209.02976, 2022.
[22] WANG C Y, BOCHKOVSKIY A, LIAO H Y M. YOLOv7: trainable bag-of-freebies sets new state-of-the-art for real-time object detectors[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023: 7464-7475.
[23] TAN M, PANG R, LE Q V. Efficientdet: scalable and efficient object detection[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2020: 10781-10790.
[24] LV W, XU S, ZHAO Y, et al. Detrs beat YOLOs on real-time object detection[J]. arXiv:2304.08069, 2023.