ATW-YOLO: reconstructing the downsampling process and attention mechanism of yolo network for rail foreign body detection
摘要
Efficient foreign body detection algorithms on railway tracks are an important prerequisite for the safe operation of trains. So, this paper proposes an improved multi-category rail foreign body detection algorithm, aiming to improve the algorithm’s ability to handle complex information. First, the conventional convolution of the YOLOv8 backbone network is replaced with the ADown downsampling module, which improves the detection ability for objects in the far distance. Second, to enhance the understanding capability for image details, a triplet attention mechanism is introduced to deal with complex texture and edge information effectively. Finally, the traditional loss function is replaced by the WIoU(Wise-IoU) loss function to optimize the training process and improve the performance of bounding box prediction regression. The experimental results on the self-built rail foreign body intrusion dataset show that compared with the baseline algorithm YOLOv8, the ATW-YOLO (A stands for ADown, T stands for triplet attention and W stands for WIoU) algorithm improves the precision by 2.2%, the recall by 1.8%, the map50 by 1.7%, the map50:95 by 3.3%, the GFLOPs and model size are reduced by 0.5 and 0.6 M, respectively. Compared with seven mainstream detection algorithms, The ATW-YOLO algorithm achieves the highest detection accuracy, and the detection speed FPS of 144 enables it to meet the real-time detection requirements of rail foreign body in working conditions.