An Improved YOLOv8 Algorithm for Rail Surface Defect Detection

Abstract

To tackle the issues raised by detecting small targets and densely occluded targets in railroad track surface defect detection, we present an algorithm for detecting defects on railroad tracks based on the YOLOv8 model. Firstly, we enhance the model’s attention towards small and medium-sized targets by substituting replacing the original convolution with the SPD-Conv building block in the backbone network of YOLOv8n, while preserving the original network structure. Secondly, we integrate the EMA attention mechanism module into the neck component, allowing the model to leverage information from different layers of features and improve feature representation capabilities. Lastly, we substitute the original C-IOU with the Focal-SIoU loss function in YOLOv8, which adjusts the weights of positive and negative samples to penalize difficult-to-classify samples more heavily. This enhancement improves the model’s capability to accurately recognize challenging samples and ensures that the network allocates greater attention to each target instance, resulting in improved performance and effectiveness of the model. The experimental results reveal notable advancements in precision, recall, and average accuracy attained by our enhanced algorithm. Compared to the original YOLOv8n model, our enhanced algorithm demonstrates remarkable precision, recall, and average accuracy of 93.9%, 93.7%, and 94.1%, respectively. These improvements amount to 3.6%, 5.0%, and 5.7%, respectively. Notably, these enhancements are accomplished while maintaining the dimensions of the model and the parameter count. During the identification of defects on railroad track surfaces, our improved algorithm surpasses other widely used algorithms in terms of performance.

Keywords

• Rail defects detection
• deep learning
• YOLOv8
• convolution module
• attention mechanism
• loss function