SGD-YOLOv5: A Small Object Detection Model for Complex Industrial Environments

Due to the complexity of industrial environments, such as construction sites and production workshops, the objects to be detected are easily occluded and perceived as small objects, which poses certain challenges for object detection. To ensure a safe industrial environment, this study adopts YOLOv5 as the basic framework and integrates the depth-to-space convolution module to improve the model’s ability to extract feature information of small targets. Second, the global attention mechanism is incorporated into the network to enhance the global interaction information, reducing the feature information loss, and improving the model performance. Finally, to alleviate the contradiction between classification and regression tasks in object detection, the YOLOv5 head is replaced with a decoupled head to achieve better classification and accelerate model convergence. To improve data diversity and enhance model robustness, we augmented the open-source safety helmet wearing dataset (SHWD) and smoking behavior detection dataset (SBDD). We test the performance of the proposed model (SGD-YOLOv5) on the large and small object detection dataset (SODA-D) and VisDrone2021-DET datasets. Furthermore, its ability to detect small objects was also evaluated. Experiments show that our model outperforms all baseline models on SHWD and SBDD. Compared to the TPH-YOLOv5 on the SODA-D dataset, AP and Recall achieve improvements of 18.3% to 19.2% and 11.9% to 13.3% respectively. On the Visdrone 2021-DET dataset, mAP@0.5 achieved an improvement from 35.45% to 35.70% compared to the state-of-the-art model YOLO-Drone.