Steel surface defect detection based on improved YOLOv5 algorithm

doi:10.11996/JG.j.2095-302X.2023020335

Abstract

Abstract:

An improved YOLOv5 steel surface defects detection algorithm was proposed to address the one-stage detection network YOLOv5, such as inadequate feature extraction ability, limited receptive field, and insufficient feature fusion. A feature pyramid structure of SPP_Res with residual edges was proposed to speed up the training of the model and enhance the feature extraction ability of the model. Additionally, a multi-head self-attention mechanism (C3_MHSA) was added to optimize the network structure, focusing on the global receptive field of the model and extracting richer features of the target. Furthermore, a multi-layer fusion mechanism was introduced to further integrate shallow and deep features, taking into account more information on location, semantics, and details, thereby improving the detection accuracy of steel surface defects. The experimental results demonstrated that the improved YOLOv5 algorithm could exhibit excellent detection performance, and that the mAP on the NEU-DET datasets reached 74.1%, which was 3.4% higher than that of the original YOLOv5 algorithm, 4.0% higher than that of the YOLOX algorithm, 8.6% higher than that of YOLOv3 algorithm, and 23.4% higher than that of the SSD algorithm. The improved YOLOv5 network could detect steel surface defects more accurately than YOLOv5 with similar detection speed, while outperforming other mainstream algorithms in both accuracy and speed.

Key words: YOLOv5, SPP_Res, muti-head self-attention mechanism, muti-layer fusion mechanism, defect detection

CLC Number:

TP391

CAO Yi-qin, WU Ming-lin, XU Lu. Steel surface defect detection based on improved YOLOv5 algorithm[J]. Journal of Graphics, 2023, 44(2): 335-345.

Figures/Tables 16

References 25

[1]	MAO T Q, REN L R, YUAN F Q, et al. Defect recognition method based on HOG and SVM for drone inspection images of power transmission line[C]// 2019 International Conference on High Performance Big Data and Intelligent Systems. New York: IEEE Press, 2019: 254-257.
[2]	CHU M X, GONG R F, GAO S, et al. Steel surface defects recognition based on multi-type statistical features and enhanced twin support vector machine[J]. Chemometrics and Intelligent Laboratory Systems, 2017, 171: 140-150. DOI URL
[3]	GIRSHICK R, DONAHUE J, DARRELL T, et al. Rich feature hierarchies for accurate object detection and semantic segmentation[C]// 2014 IEEE Conference on Computer Vision and Pattern Recognition. New York: IEEE Press, 2014: 580-587.
[4]	GIRSHICK R. Fast R-CNN[C]// 2015 IEEE International Conference on Computer Vision. New York: IEEE Press, 2016: 1440-1448.
[5]	REN S Q, HE K M, GIRSHICK R, 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. DOI PMID
[6]	REDMON J, DIVVALA S, GIRSHICK R, et al. You only look once: unified, real-time object detection[C]// 2016 IEEE Conference on Computer Vision and Pattern Recognition. New York: IEEE Press, 2016: 779-788.
[7]	REDMON J, FARHADI A. YOLO9000: better, faster, stronger[C]// 2017 IEEE Conference on Computer Vision and Pattern Recognition. New York: IEEE Press, 2017: 6517-6525.
[8]	REDMON J, FARHADI A. YOLOv3: an incremental improvement[EB/OL]. [2022-05-17]. https://arxiv.org/abs/1804.02767.
[9]	FU G Z, SUN P Z, ZHU W B, et al. A deep-learning-based approach for fast and robust steel surface defects classification[J]. Optics and Lasers in Engineering, 2019, 121: 397-405. DOI URL
[10]	LV X M, DUAN F J, JIANG J J, et al. Deep metallic surface defect detection: the new benchmark and detection network[J]. Sensors: Basel, Switzerland, 2020, 20(6): 1562.
[11]	VANNOCCI M, RITACCO A, CASTELLANO A, et al. Flatness defect detection and classification in hot rolled steel strips using convolutional neural networks[M]// Advances in Computational Intelligence. Cham: Springer International Publishing, 2019: 220-234.
[12]	HAN C J, LI G Y, LIU Z. Two-stage edge reuse network for salient object detection of strip steel surface defects[J]. IEEE Transactions on Instrumentation and Measurement, 2022, 71: 1-12.
[13]	王海云, 王剑平, 罗付华. 融合多层次特征Faster R-CNN的金属板带材表面缺陷检测研究[J]. 机械科学与技术, 2021, 40(2) : 262-269.
	WANG H Y, WANG J P, LUO F H. Study on surface defect detection of metal sheet and strip using faster R-CNN with multilevel feature[J]. Mechanical Science and Technology for Aerospace Engineering, 2021, 40(2): 262-269. (in Chinese)
[14]	代小红, 陈华江, 朱超平. 一种基于改进Faster RCNN的金属材料工件表面缺陷检测与实现研究[J]. 表面技术, 2020, 49(10): 362-371.
	DAI X H, CHEN H J, ZHU C P. Surface defect detection and realization of metal workpiece based on improved faster RCNN[J]. Surface Technology, 2020, 49(10): 362-371. (in Chinese)
[15]	LI J Y, SU Z F, GENG J H, et al. Real-time detection of steel strip surface defects based on improved YOLO detection network[J]. IFAC-PapersOnLine, 2018, 51(21): 76-81.
[16]	程婧怡, 段先华, 朱伟. 改进YOLOv3的金属表面缺陷检测研究[J]. 计算机工程与应用, 2021, 57(19): 252-258. DOI
	CHENG J Y, DUAN X H, ZHU W. Research on metal surface defect detection by improved YOLOv3[J]. Computer Engineering and Applications, 2021, 57(19): 252-258. (in Chinese) DOI
[17]	KOU X P, LIU S J, CHENG K Q, et al. Development of a YOLO-V3-based model for detecting defects on steel strip surface[J]. Measurement, 2021, 182: 109454. DOI URL
[18]	王杨, 曹铁勇, 杨吉斌, 等. 基于YOLO v5算法的迷彩伪装目标检测技术研究[J]. 计算机科学, 2021, 48(10): 226-232. DOI
	WANG Y, CAO T Y, YANG J B, et al. Camouflaged object detection based on improved YOLO v5 algorithm[J]. Computer Science, 2021, 48(10): 226-232. (in Chinese) DOI
[19]	LIN T Y, DOLLÁR P, GIRSHICK R, et al. Feature pyramid networks for object detection[C]// 2017 IEEE Conference on Computer Vision and Pattern Recognition. New York: IEEE Press, 2017: 936-944.
[20]	LIU S, QI L, QIN H F, et al. Path aggregation network for instance segmentation[C]// 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition. New York: IEEE Press, 2018: 8759-8768.
[21]	HE K M, ZHANG X Y, REN S Q, et al. Deep residual learning for image recognition[C]// 2016 IEEE Conference on Computer Vision and Pattern Recognition. New York: IEEE Press, 2016: 770-778.
[22]	RAMACHANDRAN P, PARMAR N, VASWANI A, et al. Stand-alone self-attention in vision models[EB/OL]. [2022-05-17]. 1906. 05909. https://arxiv.org/abs/1906.05909.
[23]	ZHAO H S, JIA J Y, KOLTUN V. Exploring self-attention for image recognition[C]// 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition. New York: IEEE Press, 2020: 10073-10082.
[24]	SRINIVAS A, LIN T Y, PARMAR N, et al. Bottleneck transformers for visual recognition[C]// 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition. New York: IEEE Press, 2021: 16514-16524.
[25]	XIE E Z, WANG W H, YU Z D, et al. SegFormer: simple and efficient design for semantic segmentation with transformers[J]. Advances in Neural Information Processing Systems, 2021, 34: 12077-12090.

类别	个数
CR	479
IN	741
PA	649
PS	292
RS	435
SC	395

类别	个数
CR	479
IN	741
PA	649
PS	292
RS	435
SC	395

名称	参数
GPU	RTX3060Ti-12 G
CPU	Intel(R) Core(TM) i7-10875H CPU @2.30 GHz
操作系统	Windows10
深度学习框架	Pytorch 1.9.1+cuda10.2
编译软件	PyCharm

名称	参数
GPU	RTX3060Ti-12 G
CPU	Intel(R) Core(TM) i7-10875H CPU @2.30 GHz
操作系统	Windows10
深度学习框架	Pytorch 1.9.1+cuda10.2
编译软件	PyCharm

算法	mAP	FPS	Params (M)
SSD	0.507	63.30	24.4
Cascade R-CNN	0.596	37.00	107.0
RetinaNet	0.617	42.85	28.5
YOLOv3	0.655	55.00	63.0
文献[16]	0.676	51.60	-
YOLOX(s)	0.695	102.00	9.0
YOLOX(m)	0.701	87.90	25.3
YOLOv5(m)	0.707	75.10	21.2
本文	0.741	75.00	23.9
YOLOv6(s)	0.706	121.00	17.2
YOLOv7(tiny)	0.735	165.00	6.2
YOLOv7	0.768	138.00	36.9