An improved YOLOv5 approach for real-time steel surface defect detection using attention mechanism and lightweight architecture

Abstract

Steel surface defect detection has important practical significance in industrial production. High-precision detection ensures product quality, while extreme real-time performance matches the pace of high-speed assembly lines. Existing models struggle with complex background textures and real-time edge deployment. Deployed on embedded platforms, the standard YOLOv5s faces bottlenecks in parameter count and computational burden, and is highly prone to missed detections of tiny defects under complex metal backgrounds. This paper proposes an improved YOLOv5s solution integrating an attention mechanism and a lightweight architecture. The YOLOv5s-CBAM model introduces the Convolutional Block Attention Module (CBAM), effectively suppressing background noise and significantly enhancing the feature response of weak defects. The YOLOv5s-Ghost model uses the Ghost module to reconstruct the backbone, replacing redundant feature map generation with extremely low linear computational cost to greatly reduce parameter count and FLOPs. To address the slow regression convergence of extreme aspect ratio defects like scratches, we replace the CIoU loss with SIoU, introducing joint geometric penalty constraints of angle, distance, and shape. Extensive experiments on the NEU-DET dataset show the improved lightweight model achieves an mAP@0.5 of 0.7010. The inference frame rate rises to 103.1 FPS. This effectively balances detection accuracy and edge-side computational cost, providing a reliable real-time visual inspection solution for high-speed industrial production lines.