Rail fastener detection of heavy railway based on deep learning

Abstract

Image detection based on machine learning and deep learning currently has a good application prospect for railway fault diagnosis, with good performance in feature extraction and the accuracy of image localization and good classification results. To improve the speed of locating small target objects of fasteners, the YOLOv5 framework model with faster algorithm speed is selected. To improve the classification accuracy of fasteners, YOLOv5-based heavy-duty railway rail fastener detection is proposed. The anchor size is modified on the original basis to improve the attention to small targets of fasteners. The CBAM (Convolutional Block Attention Module) module and TPH (Transformer Prediction Head) module are introduced to improve the speed and accuracy issues. The rail fasteners are divided into 6 categories. Experiment comparisons show that before the improvement, the MAP@ 0.5 value of all categories are close to the peak of 0.989 after the epoch of 150, and the F1 score approaches 1 with confidence in the interval (0.2, 0.95). The improved mAP@ 0.5 value approached the highest value of 0.991 after the epoch of 75, and the F1 score approached 1 with confidence in the interval (0.01, 0.95). The experiment results indicate that the improved YOLOv5 model proposed in this paper is more suitable for the task of detecting rail fasteners.