2.1

Detection model combining MobileNetV3 and Transformer

By integrating the parallel design of MobileNetV3 and Transformer, our model seamlessly fuses local and global features, significantly bolstering its performance. As illustrated in Figure 1, the framework comprises a backbone feature network, SPPF module, and YOLO Head. Specifically, the backbone network, dubbed MobileNetV3-CATr, integrates MobileNetV3, CA-Bneck, and Transformer modules to extract crucial rail defect features. Initially, the raw rail defect image is fed into MobileNetV3-CATr for feature extraction. Next, the SPPF module leverages varying max pooling kernels to process the defect feature maps, further enhancing feature fusion and addressing multi-scale defects. Finally, a BiFPN-Lite module is employed, and the YOLO Head precisely outputs defect details like category, location, and confidence level.

Figure 1.

Overall Diagram of the Detection Model

2.2

MobileNetV3-CATr

The MobileNetV3-CATr integrates MobileNetV3, the CA-Benck module, and the Transformer encoding module. By introducing the CA-Benck module to replace the original SE module, MobileNetV3 can more precisely focus on the damaged areas on the rail surface during feature extraction, significantly improving the accuracy of surface defect detection. Furthermore, to enhance the feature extraction capability, a Transformer encoding module is added to the end of MobileNetV3 to strengthen the capture of global information, further improving the model’s ability to recognize small defect features.

2.2.1

MobileNetV3-CA

MobileNetv3’s Bneck module with SE enhances performance but lacks positional info. CA-Bneck replaces SE with CA, considering channel & positional info. This reduces info loss, captures spatial dependencies, and precisely locates defects, improving damage detection accuracy[6].

Figure 2 displays the CA module integrating horizontal & vertical features to generate direction-aware maps. By decoupling pooling into horizontal & vertical, it encodes features independently, creating a direction-sensitive map preserving positional data. This mimics X AVG and Y AVG pooling in the CA Block. Spatial dimensions are merged, compressed via 1×1 convolutions, and fused using h-sigmoid activation, ensuring output maps retain directional sensitivity and match input channels. The module’s output is:

Figure 2.

Structure of CA-Bneck

In the equation, represents the size of the input feature map, , represent the weights for the two spatial directions respectively.

2.2.2

MobileNetv3-CA-Transformer

The images of surface defects on railway rails often suffer from sparsity issues, leading to poor detection performance for subtle defects with high similarity. By introducing the Transformer module, we enhance the model’s ability to capture global information while paying closer attention to the defect areas on the rail surface. This allows the model to further learn from regions with weaker defect features, thereby reducing the occurrence of misdetections and missed detections of surface defects on railway rails.

Figure 3.

Structure Diagram of the Transformer Module

In the context of rail surface defect detection, the self-attention and multi-head attention mechanisms of the Transformer ensure that the model focuses on defect information. By integrating the Transformer module at the end of MobileNetv3-CA, we leverage low-resolution feature maps to reduce computational and storage costs. This not only improves the mean average precision (mAP) but also reduces the network size, enabling it to excel in detecting dense and large defects, effectively enhancing the efficiency of rail surface defect detection. As shown in Figure 4, the Transformer module comprises two sublayers: a multi-head attention layer and a multi-layer perceptron (MLP). The former enhances attention to the current pixel and its contextual semantics, while the latter serves as a fully connected layer providing the functionality of a feedforward neural network. The module also includes LayerNorm and Dropout layers to enhance network integration and prevent overfitting[7].

Figure 4.

Diagram of the BiFPN-Lite Network Structure

To utilize the Transformer module, this paper flattens backbone feature maps and applies linear positional encoding. Multi-Head Self-Attention (MHSA) dynamically aggregates information via Q, K, V interactions[8], concatenating multiple heads as:

In the equation: Concat represents a tensor-level operation; X^out is a linear transformation matrix; H_m indicates that the result of the m-th self-attention is obtained through the Scaled dot-product attention in MHSA.

Finally, after the MHSA operation, two additional linear transformations are performed to obtain a feature map that focuses more on defects.

2.3

BiFPN-Lite module

The BiFPN-Lite module leverages top-down and bottom-up pathways to capture high-level semantic and low-level location information, enhancing multi-scale defect detection on rail surfaces. As seen in Figure 4, the SPPF module trains a 13×13 feature map with multi-scale techniques, resulting in 52×52 and 26×26 scales. BiFPN-Lite equally weights and normalizes features of these sizes during concatenation. The fused 13×13, 26×26, and 52×52 feature layers are then fed into the YOLO Head. These multi-scale detection heads decode to produce precise prediction boxes.

3.1

Data collection and experimental environment

An experimental site selected a 10KM railway track segment and used an intelligent inspection vehicle to capture numerous high-quality images of track surface defects. Figure 5 displays an original image of the defects.

Figure 5.

Original Image of Rail Surface Defects

The hardware configuration used in this experiment comprises a 13th Gen Intel(R) Core(TM) i7-13700 processor, an NVIDIA RTX A4000 graphics card, and 16GB of RAM. The software environment consists of Windows 11 operating system, Python 3.8 as the programming language, CUDA version 13.0, and the deep learning framework is Pytorch 1.12.1. The training iteration count is set to 200, with image Mosaic augmentation employed and its coefficient set to 0.5. Distinct from image classification, surface defect detection requires not only predicting the correct type of target but also locating its position information. The following metrics are used to evaluate the performance of defect detection:

3.2

Analysis of experimental results

To validate the detection performance of the proposed method for rail surface defects, comparative experiments were conducted on a self-constructed dataset with the improved schemes outlined in Table 1. According to the experimental results in Table 1, the introduction of the Transformer module effectively extracts contextual defect information from the rail surface, reducing the occurrence of misdetections and missed detections for some defects and enhancing the overall detection performance of the model. The CA-Bneck module enables the model to focus on both spatial and channel information of defects, allowing the backbone feature extraction network to better fit the correlation between spatial and channel information of defects, thereby facilitating the extraction of feature information for rail surface defects. The mAP value of the proposed model in this paper increased by 9.5% compared to the original model, indicating that MobileNetV3-CATr and BiFPN-Lite have a positive effect on the extraction of defect features on the rail surface. This demonstrates the positive interaction among various modules and their ability to improve detection accuracy.

Table 1.

Comparative Experiments Among Different Models

To further emphasize the effectiveness of our model, we’ve visualized select experimental results. As depicted in Table 2, our proposed algorithm rarely misses defects. Unlike the original model, which is limited to detecting large, clear targets, our enhanced algorithm adeptly detects even blurry and dark defects. Moreover, in cases where two defects overlap, our algorithm can distinguish and precisely locate the two nearby targets, including smaller ones within larger ones, thereby greatly enhancing defect detection accuracy and minimizing missed detections.

Table 2.

Visualization Results of Different Schemes

3.3

Experimental comparison of different algorithms

To validate the performance of the proposed method for railway rail surface defect detection, this paper conducts a comparative analysis with other detection methods. As shown in Table 3, the proposed algorithm demonstrates superior performance overall. Compared to mainstream one-stage object detection algorithms such as YOLOv3, YOLOv4, and YOLOv5, it requires the least number of parameters (Params) and achieves higher accuracy. In terms of detection speed, the proposed algorithm achieves 20.1 FPS on devices with lower configurations, which can basically meet the requirements of online real-time detection.

Table 3.

Comparison Results of Different Algorithms