计算机科学 ›› 2023, Vol. 50 ›› Issue (11A): 230900113-7.doi: 10.11896/jsjkx.230900113

• 图像处理&多媒体技术 • 上一篇    

改进YOLOv5s的轻量化钢材表面缺陷检测模型

蒋博1, 万毅1, 谢显中2   

  1. 1 重庆邮电大学通信与信息工程学院 重庆 400065
    2 重庆邮电大学研究生院 重庆 400065
  • 发布日期:2023-11-09
  • 通讯作者: 蒋博(2905828080@qq.com)

Improved YOLOv5s Lightweight Steel Surface Defect Detection Model

JIANG Bo1, WAN Yi1, XIE Xianzhong2   

  1. 1 School of Communication and Information Engineering,Chongqing University of Posts and Telecommunications,Chongqing 400065,China
    2 Graduate School,Chongqing University of Posts and Telecommunications,Chongqing 400065,China
  • Published:2023-11-09
  • About author:JIANG Bo,born in 1998,postgraduate.His main research interests include machine learning and image processing.

摘要: 针对现有钢材表面缺陷检测模型结构复杂、参数量多、检测精度和实时性较差等问题,提出了一种改进YOLOv5s的轻量化钢材表面缺陷检测模型。首先采用MobileNetv3-Small网络替换YOLOv5s主干提取网络,实现模型轻量化,提升检测速度;其次在特征融合阶段采用加权双向特征金字塔网络(Bidirectional Feature Pyramid Network,BiFPN)加强特征提取,通过融合不同尺度的特征,提升检测的准确率和鲁棒性。同时引入CBAM(Convolutional Block Attention Module)注意力机制增强模型对小尺度目标的检测能力;最后使用K-means++算法聚类先验框,提高先验框聚类的准确性和收敛速度。改进后的模型在NEU-DET数据集上的平均精度均值(mAP@0.5)达到77.2%,在NVIDIA 1080Ti上检测速度达到102FPS。相较于原始YOLOv5s模型,mAP提升3.90%,参数量减少58.6%,体积减小34%,检测速度提升29.7%。实验结果表明改进的YOLOv5s模型在保证轻量化的同时能够有效提升钢材表面缺陷检测的精度和速度,易于部署,满足带钢实际生产中的需求。

关键词: 缺陷检测, YOLOv5s, 轻量化, MobileNetv3-Small, BiFPN, CBAM, K-means++

Abstract: Aiming at the problems of complex structure,large number of parameters,poor detection accuracy and real-time performance of existing steel surface defect detection models,this paper proposes an improved YOLOv5s lightweight steel surface defect detection model.Firstly,the MobileNetv3-Small is used to replace the YOLOv5s backbone extraction network,achieving model lightweight and improving detection speed.Secondly,in the feature fusion stage,a weighted bidirectional feature pyramid network (BiFPN) is used to enhance feature extraction.By fusing features of different scales,the accuracy and robustness of detection are improved.Simultaneously,the convolutional block attention module(CBAM) attention mechanism is introduced to enhance the model's ability to detect small scale targets.Finally,the K-means++ algorithm is proposed to cluster prior boxes,improve the accuracy and convergence speed of prior box clustering.The average accuracy of the improved YOLOv5s on the NEU-DET dataset (mAP@0.5) reaches 77.2%,with a detection speed of 102 FPS on NVIDIA 1080Ti.Compared to the original YOLOv5s,the mAP is increased by 3.90%,the parameter quantity is decreased by 58.6%,the volume is decreased by 34%,and the detection speed is increased by 29.7%.Experimental results demonstrate that the improved lightweight YOLOv5s effectively improves both the accuracy and speed of steel surface defect detection.Moreover,it is easy to deploy and meet the requirements of actual production in the steel strip industry.

Key words: Defect detection, YOLOv5s, lightweight, MobileNetv3-Small, BiFPN, CBAM, K-means++

中图分类号: 

  • TP391
[1]ZHANG R,WANG Y Y,DUAN Y Q,et al.Real time object detection and localization algorithm for flight robotic arms[J].Journal of Nanjing University of Aeronautics and Astronautics,2022,54 (1):27-33.
[2]GIRSHICK R.Fast r-cnn[C]//Proceedings of the IEEE International Conference on Computer Vision.2015:1440-1448.
[3]DAI J,LI Y,HE K,et al.R-fcn:Object detection via region-based fully convolutional networks[J].Advances in Neural Information Processing Systems,2016,29:379-387.
[4]HE K,GKIOXARI G,DOLLÁR P,et al.Mask r-cnn[C]//Proceedings of the IEEE International Conference on Computer Vision.2017:2961-2969.
[5]LIU W,ANGUELOV D,ERHAN D,et al.Ssd:Single shotmultibox detector[C]//Computer Vision-ECCV 2016:14th European Conference,Amsterdam,The Netherlands(Part I 14).Springer International Publishing,2016:21-37.
[6]REDMON J,DIVVALA S,GIRSHICKR,et al.You only lookonce:Unified,real-time object detection[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition.2016:779-788.
[7]REDMON J,FARHADI A.Yolov3:An incremental improve-ment[J].arXiv:1804.02767,2018.
[8]BOCHKOVSKIY A,WANG C Y,LIAO H Y M.Yolov4:Optimal speed and accuracy of object detection[J].arXiv:2004.10934,2020.
[9]GE Z,LIU S,WANG F,et al.Yolox:Exceeding yolo series in 2021[J].arXiv:2107.08430,2021.
[10]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.
[11]YANG J H,LI H,DU Y Y,et al.Lightweight object detection algorithm based on improved YOLOv5s[J].Electro Optics and Control,2023,30 (2):24-30.
[12]CAO Y Q,WU M L,XU L.Steel surface defect detection based on improved YOLOv5 algorithm[J].Journal of Graphics,2023,44 (2):335
[13]LI X,WANG C,LI B,et al.Improvement of YOLOv5 steel surface defect detection algorithm [J].Journal of Air Force Engineering University(Natural Science Edition),2022,23(2):26-33.
[14]WU D,LI M H,MA W K,et al.Steel surface defect detection based on improved YOLOv5[J].Journal of Shanxi University of Science&Technology,2023,41 (2):162-169.
[15]HOWARD A,SANDLER M,CHU G,et al.Searching for MobileNetv3[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision.2019:1314-1324.
[16]HOWARD A G,ZHU M,CHEN B,et al.Mobilenets:Efficient convolutional neural networks for mobile vision applications[J].arXiv:1704.04861,2017.
[17]SANDLER M,HOWARD A,ZHU M,et al.Mobilenetv2:Inverted residuals and linear bottlenecks[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition.2018:4510-4520.
[18]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.
[19]WOO S,PARK J,LEE J Y,et al.Cbam:Convolutional block attention module[C]//Proceedings of the European Conference on Computer Vision(ECCV).2018:3-19.
[20]MODHA D S,SPANGLER W S.Feature weighting in k-means clustering[J].Machine learning,2003,52:217-237.
[21]ARTHUR D,VASSILVITSKII S.K-means++ the advantages of careful seeding[C]//Proceedings of the Eighteenth Aannual ACM-SIAM Symposium on Discrete Algorithms.2007:1027-1035.
[22]HE Y,SONG K,MENG Q,et al.An end-to-end steel surfacedefect detection approach via fusing multiple hierarchical features[J].IEEE Transactions on Instrumentation and Measurement,2019,69(4):1493-1504.
[23]ZHANG X,ZHOU X,LIN M,et al.Shufflenet:An extremelyefficient convolutional neural network for mobile devices[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition.2018:6848-6856.
[24]HAN K,WANG Y,TIAN Q,et al.Ghostnet:More featuresfrom cheap operations[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition.2020:1580-1589.
Viewed
Full text


Abstract

Cited

  Shared   
  Discussed   
No Suggested Reading articles found!