Improved YOLOv5 Network for Steel Surface Defect Detection

Abstract

Steel surface defect detection is crucial for ensuring steel quality. The traditional detection algorithm has low detection probability. This paper proposes an improved algorithm based on the YOLOv5 model to enhance detection probability. Firstly, deformable convolution is introduced in the backbone network, and a traditional convolution module is replaced by deformable convolution; secondly, the CBAM attention mechanism is added to the backbone network; then, Focal EIOU is used instead of the CIOU loss function in YOLOv5; lastly, the K-means algorithm is used to cluster the Anchor box, and the Anchor box parameters that are more suitable for this paper are obtained. The experimental results show that using deformable convolution instead of traditional convolution can get more feature information, which is more conducive to the learning of the network. This paper uses the CBAM attention mechanism, and the heat map of the attention mechanism shows that the CBAM attention mechanism is beneficial for feature extraction. Focal EIOU is optimized in high and wide loss compared with the CIOU loss function, which accelerates the convergence of the model. The Anchor box is more favorable for feature extraction. The improved algorithm achieved a detection probability of 78.8% in the NEU-DET dataset, which is 4.3% better than the original YOLOv5 network, and the inference time of each image is only increased by 1 ms; therefore, the optimized algorithm proposed in this paper is effective.

Keywords

• YOLOv5
• deformable convolution
• attention mechanism
• <i>Focal EIOU</i>
• <i>K-means</i>