SSA-YOLO: An Improved YOLO for Hot-Rolled Strip Steel Surface Defect Detection

Abstract

In the manufacturing process of hot-rolled steel strips, various mechanical forces, and environmental conditions can cause surface defects, making their detection crucial for ensuring high-quality product production and preventing significant economic losses in the industry. However, existing models within the you only look once (YOLO) family, commonly employed for steel surface defect detection, have exhibited limited effectiveness. In this article, we propose an improved version of YOLO, namely, YOLO enhanced by a convolution squeeze-and-excitation (CSE) module, Conv2d-BatchNorm-SiLU (CBS) with Swin transformer (CST) module, and adaptive spatial feature fusion (ASFF) detection head module, i.e., SSA-YOLO, specifically tailored for end-to-end surface defect detection. Our approach incorporates several key modifications aimed at improving performance. First, we integrate a channel attention mechanism module into the shallow convolutional network module of the backbone. This enhancement focuses on channel information to improve feature extraction related to small defects while reducing redundant information in candidate boxes. In addition, we fuse a Swin transformer (Swin-T) module into the neck to enhance feature representation for detecting diverse and multiscale defects. Finally, the ASFF is introduced in YOLO to increase cross-interaction between high and low levels in the feature pyramid network (FPN). Experimental results demonstrate the superior performance and effectiveness of our SSA-YOLO model compared to other state-of-the-art models. Our approach achieves higher accuracy and sensitivity in detecting surface defects, offering significant advancements in steel strip production quality control. The code is available at https://github.com/MIPIT-Team/SSA-YOLO .