Lightweight Detection Method of Wheelset-Tread Defects Based on Improved YOLOv7

1. Introduction

Train wheelsets, which are critical components of high-speed trains, are susceptible to failure under long-term operation in harsh environments [1,2]. The defects in wheelset treads typically include peels, wear, bruises, and pits. The efficient detection of wheelset-tread defects is necessary for safe operation of trains [3,4]. Traditional methods primarily adopt manual measurements; however, most inspection information is not used appropriately for efficient detection [5]. Machine vision technology is frequently employed in tedious inspection tasks to reduce labor waste [6,7,8]. In complex environments, traditional machine vision inspection faces challenges such as vibration, noise, and light variation, which seriously influence measurement accuracy and efficiency. Numerous defect detection studies have focused on convolutional neural networks (CNNs) because of their superior ability to overcome various disturbances [9]. The most common models mainly include two-stage models, such as faster region-based CNN (Faster R-CNN) models, for their accuracy [10,11], and two-stage models, such as You Only Look Once (YOLO), for their efficiency [12,13]. YOLO-based algorithms have been widely applied in numerous online and complex detection areas owing to their high efficiency [14,15,16,17].

Accurate and comprehensive online inspection of defects in wheelset treads are difficult because of their variety and position uncertainty; in particular, the accurate detection and location of small targets poses greater challenges. Compared to common targets, small targets usually have a low pixel count with a smaller portion and are easily overlapped; they are therefore difficult to recognize. Ren et al. [18] improved YOLOv4 to solve the problem of small target recognition by using a lightweight backbone and a shallow feature enhancement network. Tian et al. [19] proposed multiscale dense YOLO (MD-YOLO), which incorporates a feature extraction path and feature aggregation path. A deep network, which leverages spatial location information from a shallower network, results in enhanced accuracy of small target detection. Wei et al. [20] modified YOLOv3 model to rectify the shortcomings of small target detection. Huang et al. [21] addressed the problem of poor small target detection and optimized k-means clustering algorithms with an added small target detection layer. Lin et al. [22] reported an enhancement in target localization ability of a network with the use of the coordinate attention (CA) mechanism within the feature aggregation process design. Cai et al. [23] achieved a more precise and efficient detection method that utilizes channel attention and spatial attention modules to suppress insignificant features in datasets. Kang et al. [24] proposed ASF-YOLO, which combines spatial and scale features for accurate and fast cell-instance segmentation. Li et al. [25] proposed a new lightweight convolution technique (GSConv) to reduce the model size while maintaining accuracy.

This paper proposed an improved, YOLOv7-based model that realizes efficient detection of defects in the wheelset tread of high-speed trains. The captured images of the wheelset tread are preprocessed for small and nonuniform samples. First, the image dataset of the wheelset tread are expanded using traditional data augmentation and Copy–Paste [26] to address data scarcity. Second, the new dataset is trained using the StyleGAN3 network [27]. Third, the defects are labeled using the LabelImg software. Then, the dataset is divided into training, validation, and test sets. Finally, lightweight convolution in the Neck module is used to reduce the computational parameters of the model, and a small target enhancement (STE) module is added to address the issue of small sample detection. A comparison of the experimental results indicates that YOLOv7-STE exhibits better detection accuracy. A publicly available rail surface defect dataset (RSDDs) dataset with similar small target defects was used to verify the robustness of the proposed model. As compared to the traditional network model, the proposed model shows good performance in detecting small target defects in different samples. The experimental results confirm that YOLOv7-STE can be applied for the online inspection of wheel tread defects.

2. Materials and Methods

2.1. Dataset, Environment, and Parameters

The representational ability of deep models is highly dependent on the diversity of training data. However, fully publicly available wheelset damage target datasets are not available. Therefore, deep learning-based wheelset damage-detection methods necessitates the creation and expansion of datasets.

The required images of various types of defects are captured in the field; however, the number of useful defect images obtained is limited and typically unbalanced. The acquired data images generally present three types of defects, namely peels, pits, and bruises, as shown in Figure 1. Given the insufficient volume of the dataset, it was preprocessed to expand the samples. It was first expanded using traditional data augmentation and Copy–Paste; the results are shown in Figure 2. A dataset comprising 654 images was obtained; however, it was insufficient. Subsequently, StyleGAN3 network was trained. Figure 3 shows a generative image of wheelset-tread defects by using StyleGAN3. To address this imbalance in the dataset, StyleGAN3 was used to selectively generate the corresponding defective images for dataset expansion, and a new dataset comprising 1200 images was obtained. The defects were labelled using LabelImg software, and the dataset was divided into training, validation, and test sets in an 8, 1, 1 ratio. The training set was further expanded to 2160 images by splicing and blending the augmented data with Gaussian blur, affine transform, luminance transform, descending pixel transform, and flip transform [28,29]. The statistical information for the dataset is presented in Table 1. Figure 4 shows significant dimensional differences among the different defects.

The hardware environment and software versions of the experiments are listed in Table 2. In this study, stochastic gradient descent (SGD) method was used to optimize the learning rate, and the epochs were determined by comparing the loss functions of training and validation sets. The parameters of the training network are listed in Table 3.

2.2. Loss Function and Model Evaluation Metrics

To verify the superior performance of improved YOLOv7 model, we determined the mean average precision (mAP), frames per second (FPS), and model volume. The commonly used metrics of precision (P), recall (R), and mAP were selected to evaluate model performance [30]; these metrics are defined as follows:

$$P_{mAP}={\displaystyle \frac{1}{C_{class}}}{\displaystyle \sum _{i=1}^{C_{class}}{\displaystyle \underset{0}{\overset{1}{\int }}P(R)dR}}$$

(1)

$$P={\displaystyle \frac{N_{TP}}{N_{TP}+N_{F\mathrm{P}}}}$$

(2)

$$R={\displaystyle \frac{N_{TP}}{N_{TP}+N_{F\mathrm{N}}}}$$

(3)

Here, TP denotes positive samples predicted to be correct (i.e., true positives), FP denotes negative samples predicted to be incorrect (i.e., false positives), FN denotes positive samples predicted to be incorrect (i.e., false negatives), and N denotes the number of sample categories.

The loss function of YOLOv7 comprises the following three components: confidence, bounding-box regression, and classification losses. The YOLOv7 loss function can be expressed as follows:

$${Loss}_{total}={\mathsf{\lambda }}_1{\mathrm{L}}_{\mathrm{o}\mathrm{b}\mathrm{j}}+{\mathsf{\lambda }}_2{\mathrm{L}}_{box}+{\mathsf{\lambda }}_3{\mathrm{L}}_{cls}$$

(4)

${\mathrm{L}}_{\mathrm{o}\mathrm{b}\mathrm{j}}$, ${\mathrm{L}}_{box}$, and ${\mathrm{L}}_{cls}$ denote confidence, bounding-box regression, and classification losses, respectively. Furthermore, ${\mathsf{\lambda }}_1$, ${\mathsf{\lambda }}_2$, and ${\mathsf{\lambda }}_3$ are weight coefficients for the three losses; changing these coefficients can adjust the emphasis on the three losses. In YOLOv7, ${\mathrm{L}}_{box}$ is calculated using the complete intersection over union (CIoU) loss, ${\mathrm{L}}_{\mathrm{C}\mathrm{I}\mathrm{o}\mathrm{U}}$[31], which improves the speed and accuracy of bounding box regression. The expression for ${\mathrm{L}}_{\mathrm{C}\mathrm{I}\mathrm{o}\mathrm{U}}$ is as follows:

$$L_{CIOU}=1-IOU+{\displaystyle \frac{{\rho }^2\left(b,b_{gt}\right)}{c^2}}+\alpha \nu$$

(5)

$$IOU={\displaystyle \frac{\left|b\cap b_{gt}\right|}{\left|b\cup b_{gt}\right|}}$$

(6)

$$\nu ={\displaystyle \frac{4}{{\pi }^2}}{\left(\arctan {\displaystyle \frac{w_{gt}}{h_{gt}}}-\arctan {\displaystyle \frac{w}{h}}\right)}^2$$

(7)

Here, $b$ and $b_{gt}$ denote the predicted box and ground truth box, respectively; $w_{gt}$, $h_{gt}$, $w$, and $h$ denote the width and height of the ground truth box and predicted box, respectively; $\rho$ represents the distance between the centers of the two boxes; $c$ represents the maximum distance between the boundaries of the two boxes; and $\alpha$ denotes a weight coefficient.

2.3. Improved YOLOv7 Network Architecture

2.3.1. YOLOv7 Network Architecture

YOLOv7 is a widely used detection network. The network structure comprises Head, Neck, and Backbone modules as shown in Figure 5. YOLOv7 typically comprises three common models: YOLOv7, YOLOv7x, and YOLOv7tiny. Specifically, YOLOv7tiny is the smallest model with a fast detection speed and relatively low accuracy, and it is mainly applied to scenes with high speed requirements or limited computational resources. Furthermore, YOLOv7 is suitable for scenes requiring high detection accuracy. In this study, YOLOv7 was used as the base model, and a few parameters were adjusted to improve the model performance.

2.3.2. YOLOv7-STE

The Neck module of YOLOv7 was improved to enhance the correction of the detection and classification of different types of wheelset-tread defects. GSConv was introduced instead of Conv, and it can guarantee detection accuracy, reduce the network volume and computation amount, and obtain a lightweight design of the model. To detect and recognize small pits in the defect images, an STE module was added, which performed feature fusion of three different scales of feature maps (p3, p4, and p5) in the Backbone module to capture multiscale features. The structure of the YOLOv7-STE network is illustrated in Figure 6.

The main function of the multiscale module involves the fusion of features from different scales in the Backbone module via 3D convolution and pooling; in particular, it provides better feature representation for the detection of small targets. Furthermore, channel-attention and spatial-attention modules are introduced for multiscale features to fully use cross-channel information and achieve improved performance without increasing network complexity, as shown in Figure 7. Noise and irrelevant information were suppressed, whereas critical information was retained.

2.3.3. Loss Function

The CIoU is sensitive to the aspect ratio; thus, aspect ratio matching can be overemphasized at the expense of other factors such as the overall positional relationship of the target. Enhanced intersection over union (EIoU) [32] is mainly used in target detection to evaluate the difference between predicted and real bounding boxes. The expression for $\mathrm{E}\mathrm{I}\mathrm{o}\mathrm{U}$ is presented in Equation (11).

Given that the distance from the center is $D_c$, the center of the prediction box is $\left(x_p,y_p\right)$, true box center is $\left(x_g,y_g\right)$, width difference is $w_{diff}$, and height difference is $h_{diff}$. Hence,

$$D_c=\sqrt{{\left(x_p-x_g\right)}^2+{\left(y_p-y_g\right)}^2},$$

(8)

$$w_{diff}=\left|w_p-w_g\right|,$$

(9)

$$h_{diff}=\left|h_p-h_g\right|,$$

(10)

$$EIoU=IoU-{\lambda }_1{\displaystyle \frac{D_C^2}{C_d^2}}-{\lambda }_2{\displaystyle \frac{w_{diff}^2}{C_w^2}}-{\lambda }_3{\displaystyle \frac{h_{diff}^2}{C_h^2}}$$

(11)

where${\lambda }_1$, ${\lambda }_2,$and ${\lambda }_3$ are hyperparameters that regulate the effects of center distance and width-height differences; and $C_d$, $C_w$, and $C_h$ are constants used for normalization, and they are related to the size of the target in the dataset.

The utilization of EIoU provides several advantages. First, the comprehensive consideration of multiple geometric information of the bounding box, including overlapping area, center distance, and width–height difference, enables the model to more accurately learn the position and shape information of the target. This improves the accuracy of target detection, particularly when in complex situations such as deformation and occlusion of the target. Second, compared to some traditional loss functions, the width and height loss of EIoU directly minimizes the difference between the width and height of the target and prediction frames; this can guide the model to quickly learn the appropriate bounding box parameters, accelerating the convergence speed of the model. Third, the EIoU loss function enhances the small target bounding box.

3. Experimental Results and Discussion

3.1. Comparative Experimental Results Analysis

The experimental results of the proposed model and the mainstream detection algorithms are compared in Table 4. The mAP of the proposed model was higher by 1.6%, 10.7%, 48.63%, and 37.97% as compared to those of YOLOv7, YOLOv5, single shot multibox detector (SSD), and Faster R-CNN models, respectively. The model parameter size was lower by 73.91 MB, 94.69 MB, 122.11 MB, and 154.91 MB, respectively, when compared to those of the aforementioned detection models. The proposed model exhibited better detection accuracy than did the other models. Table 4 highlights the significantly better overall performance of the proposed model.

3.2. Ablation Experiment

The ablation experiment explored the impact of the improved methods on the network model. The plotted data are presented in Table 5 and Table 6. We conducted eight groups of experiments, adding different modules, and compared them with the original YOLOv7 model, using the mAP@0.5, model volume, mAP@0.5, 0.95, and FPS as indicators. Of the five different loss functions used on the original YOLOv7 model, EIoU loss yielded the best performance (Table 5). mAP@0.5 increased by 1.8% and mAP@0.5, 0.95 increased by 1.44% as compared to the original YOLOv7 model (with CIoU loss). The EIoU loss function therefore improves the effect of small-object bounding boxes.

Table 6 shows that the model volume of the GSConv module in the Neck module reduced by 84 MB, whereas the FPS increased to 87.1. As regards the small target (the pit) detection, mAP@0.5 reduced by 10.2%. The introduction of the STE module improved the small-target detection mAP by 4.4%, while the model volume increased by 14 MB. The detection speed decreased, and the mAP was lower for the rest of the target sizes, with a 2.6% reduction in mAP@0.5 for the peel and a 1.3% reduction in mAP@0.5 for the bruise. When both the STE module and GSConv were used, compared to YOLOv7 (EIoU), the model volume was reduced by 73.91 MB. The mAP@0.5 (all classes) improved by 1.6%, with a 4.2% improvement in the mAP of the pit, 0.1% in the mAP of the peel, 0.3% in the mAP of the bruise, and 1.9% in FPS. In summary, the improved strategy based on YOLOv7 proposed herein leads to significant improvement in for wheelset-tread damage detection.

3.3. Comparison of Different Models

To verify the superiority of the proposed YOLOv7-STE model for wheelset-tread damage detection, we compared the results with the outcomes from YOLOv7, YOLOv5, SSD, and Faster R-CNN models. The relevant parameters of the experiments were strictly controlled using a uniform image size as the input and a uniform training and test set for experimental testing. A comparison of the recognition of pits using the different models is shown in Figure 8. The confidence levels of detection results of YOLOv7-STE, YOLOv7, YOLOv5, SSD, and Faster R-CNN are 0.96, 0.86, 0.77, 0.69, and 0.58, respectively. The recognition of bruises using the different models is shown in Figure 9. The confidence levels of YOLOv7-STE, YOLOv7, and YOLOv5, and SSD, and Faster R-CNN detection results were 0.95, 0.91, 0.92, 0.85, and 0.77, respectively. Figure 10 shows a comparison of the recognition results of the peels by using different models. The confidence levels of YOLOv7-STE, YOLOv7, YOLOv5, SSD, and Faster R-CNN detection results were 0.97, 0.95, 0.95, 0.90, and 0.89, respectively. The experimental results indicate that the proposed method was more adaptable to complex backgrounds than the comparison methods and achieved better inspection performance.

A publicly available RSDDs dataset with similarly small target defects was used for verifying the robustness of the model. The experimental results are presented in Figure 11. The confidence level of YOLOv7-STE is remarkably high, at 0.97. Furthermore, the confidence levels of YOLOv7, YOLOv5, SSD, and Faster-RCNN are 0.90, 0.88, 0.52, and 0.49, respectively.

Images detection results showed that, for large-size defects (peel and bruise), the five models were effective, with YOLOv7-STE exhibiting the highest correct classification rate. For small defects (pit), YOLOv7-STE exhibited the highest recognition rate and correct classification rate. In summary, the experimental results confirm that the YOLOv7-STE model outperformed the other four models. When tested on RSDDs dataset, YOLOv7-STE demonstrated superior performances compared to traditional deep neural networks. The proposed model showed remarkable robustness to small target defects in different samples.

4. Conclusions

This study developed an improved YOLOv7 model was developed for the inspection of wheelset-tread defects, offering improved accuracy and reduced power computation. GSConv, a lightweight convolution in the Neck module, was adopted to reduce the model volume. The STE module addressed the challenge of low pixels and difficulty in distinguishing small targets, and improving the classification accuracy for damage feature recognition. The captured images of the wheelset tread were preprocessed to augment small and nonuniform samples through Copy–Paste, followed by training with the StyleGAN3 network. A comparison of the experimental results showed that the mAP of the proposed model improved by 1.6%, 10.7%, 48.63%, and 37.97% as compared to those of YOLOv7, YOLOv5, SSD, and Faster R-CNN models, respectively. The model parameter size decreased by 73.91 MB, 94.69 MB, 122.11 MB, and 154.91 MB as compared to the aforementioned detection models, respectively. The YOLOv7-STE model exhibited a significantly better overall performance than the traditional models. Its robustness was verified using a publicly available RSDDs dataset with similar small target defects. Compared to the traditional models, YOLOv7-STE showed better performance for small target defects in different samples. The experimental results confirmed that the YOLOv7-STE model can satisfy the requirements of online inspection of wheelset-tread defects.

However, in this study, owing to the limited collection of defect data, only three common types of damage (pit, bruise, and peel) were identified and classified. In real scenarios, owing to the influence of a variety of complex factors (including differences in temperature and light intensity), similar defects (or even the same defects in different spatial and temporal) can lead to differences in manifestation. Hence, it is necessary to further expand the data and conduct different types of classification and identification.