A lightweight real-time detection method and device for track surface defects

Abstract

The present invention discloses a lightweight real-time detection method and device for track surface defects, which relate to the field of image data processing technology. By optimizing the model structure and training strategy, the detection accuracy and efficiency of track defect detection are improved. The method includes: using a track image dataset to pre-train a T-YOLOv8 teacher model; the T-YOLOv8 teacher model is based on the original YOLOv8 model, introducing the NEW-IAT module in the backbone network, and replacing the neck network with the CBAM-C2f module to obtain the model; using the trained T-YOLOv8 teacher model to perform distillation training on a lightweight S-YOLOv8 student model, and optimizing and verifying the trained lightweight S-YOLOv8 student model; the lightweight S-YOLOv8 student model is based on the original YOLOv8 model, retaining the SPPF layer, and replacing the rest of the backbone network with a lightweight new backbone network, which includes three modules: HGStem, Light_HGBlock, and DWConv; and using the optimized lightweight S-YOLOv8 student model to perform real-time detection of track surface defects to obtain track defect information.

Description

Technical Field

[0001] The present invention relates to the technical field of image data processing, and in particular to a lightweight real-time detection method and device for track surface defects. Background Art

[0002] The rail transit system is mainly composed of a catenary system, track lines, vehicles, stations, etc. During long-term use, the rails on the track lines are affected by both the natural environment and train operation, and are prone to various defects such as cracks, breakages, and wear. If these defects are not discovered and handled in a timely manner, they are very likely to lead to serious accidents such as train derailment and collision, causing huge casualties and property losses, posing a huge challenge to the operational guarantee and further healthy development of rail transit.

[0003] With the development of computer vision and deep learning technologies, automated track defect detection based on image and video data has become an effective solution for track defect detection. In the specific task of railway defect detection, the wide variety of surface defects and their diverse forms make it challenging to accurately identify and classify these different types of defects. Furthermore, the complex background of the rails, with their abundant non-defect interference, poses a challenge in accurately locating and identifying tiny defects within this complex environment. These difficulties contribute to the persistent low accuracy and efficiency of existing detection technologies. Summary of the Invention

[0004] In order to overcome the above technical defects, the present invention provides a lightweight real-time detection method and device for track surface defects, which improves the detection accuracy and efficiency of track defect detection by optimizing the model structure and training strategy.

[0005] In a first aspect, the present invention provides a lightweight real-time detection method for track surface defects, comprising:

[0006] S1. Acquire track images to form a dataset, which is divided into a training set, a test set, and a validation set;

[0007] S2. Pre-train the T-YOLOv8 teacher model using the dataset;

[0008] The T-YOLOv8 teacher model is based on the original YOLOv8 model. The NEW-IAT module is introduced into the backbone network and the neck network is replaced with the CBAM-C2f module. The CBAM-C2f module includes a dilated convolutional layer, a channel attention module, and a spatial attention module.

[0009] S3. Use the trained T-YOLOv8 teacher model to perform distillation training on the lightweight S-YOLOv8 student model, and optimize and verify the trained lightweight S-YOLOv8 student model;

[0010] The lightweight S-YOLOv8 student model is based on the original YOLOv8 model, retaining the SPPF layer and replacing the original YOLOv8 backbone network with a lightweight new backbone network. The lightweight new backbone network includes three modules: HGStem, Light_HGBlock, and DWConv.

[0011] S4. Use the optimized lightweight S-YOLOv8 student model to detect track surface defects in real time to obtain and display track defect information.

[0012] Optionally, in a possible implementation of the first aspect, the NEW-IAT module includes a first branch and a second branch, the first branch includes an expanded convolution layer and a PEM module, the second branch includes an expanded convolution layer and an improved GPM module, the improved GPM module uses a multi-head attention mechanism to replace the original cross attention mechanism, wherein the PEM module is used to enhance local feature details, and the improved GPM module is used to generate global adjustment parameters; the CBAM-C2f module includes an expanded convolution layer, a channel attention module and a spatial attention module; its structure and feature extraction process are as follows: image input First, a feature map is generated by a 1x1 dilated convolution layer; then the feature map is divided into two parts through a two-branch structure: the first part is directly passed, and the second part is extracted by several 3x3 dilated convolution layers. The two parts of the feature are then spliced ​​in the channel dimension to obtain a new feature map; the new feature map after splicing enters the CBAM module, first it is weighted in the channel dimension by the channel attention module, and then it is weighted in the spatial dimension by the spatial attention module to obtain an enhanced feature map. Finally, the enhanced feature map passes through a 1x1 convolution layer to output the final feature map.

[0013] Optionally, in a possible implementation of the first aspect,

[0014] In S2, the specific steps for pre-training the T-YOLOv8 teacher model include:

[0015] S2.1. Use the NEW-IAT module to perform dark light enhancement or exposure correction on the image.

[0016] S2.2, use the CBAM-C2f module to perform feature extraction on the image processed in S2.1;

[0017] S2.3. Input the extracted image features into the prediction head in the T-YOLOv8 teacher model to generate the final prediction results of object detection.

[0018] Optionally, in a possible implementation of the first aspect,

[0019] In S3, the specific steps of distillation training include:

[0020] S3.1. Design a distillation loss function based on classification loss, localization loss, and distillation loss.

[0021] S3.2. During the distillation training process, the track image is used as input, and iterative training is performed by adjusting the weights in the distillation loss function until the number of iterations meets the threshold.

[0022] Optionally, in a possible implementation of the first aspect, after completing distillation training, optimizing and verifying the lightweight S-YOLOv8 student model includes:

[0023] S3.3. Model Optimization: Use an optimization strategy to optimize the lightweight S-YOLOv8 student model to improve the performance of the lightweight S-YOLOv8 student model. The optimization strategy is to adjust the weights of the lightweight S-YOLOv8 student model in combination with the conventional object detection loss function.

[0024] S3.4. Effect evaluation and verification: Use an independent test set to evaluate the performance of the lightweight S-YOLOv8 student model trained after distillation.

[0025] Optionally, in a possible implementation of the first aspect, the display is to present the detected track defect information through an interactive web interface; the web interface includes marking the detected defect location and type, defect pictures and detailed descriptions.

[0026] Optionally, in a possible implementation of the first aspect, track images are obtained to form a dataset, and the dataset is divided into a training set, a test set, and a validation set, including:

[0027] Railway track images are collected using cameras and other image acquisition devices. After data cleaning and labeling, the labeled track images are divided into training set, test set, and validation set according to preset ratios.

[0028] Optionally, in a possible implementation of the first aspect, the labeling tool LabelImg is used for data labeling, and LabelImg is used to perform the following image labeling process: create labels in LabelImg, and then convert the labeling information into the txt file format required for model training, where each image corresponds to a .txt file, which contains the category label of the target and the coordinate information of its bounding box.

[0029] In the second aspect, the present invention provides a lightweight real-time detection device for track surface defects, comprising: an image acquisition module, a data processing module, a communication module, a power management module and a display module, wherein the data processing module is used to execute a lightweight real-time detection method for track surface defects described in any one of the above-mentioned first aspect and its implementation methods.

[0030] In the third aspect, another aspect of the present invention provides a computer-readable storage medium, wherein the computer-readable storage medium stores computer instructions, and when the computer instructions are executed by a processor, a lightweight real-time detection method for track surface defects described in any one of the above-mentioned first aspect and its implementation methods is implemented.

[0031] In a fourth aspect, another aspect of the present invention provides a computer program product, comprising a computer program, which, when executed by a processor, implements a lightweight real-time detection method for track surface defects as described in any one of the first aspect and its implementation methods.

[0032] The above-mentioned technical solution employed in the present invention achieves the following technical effects: Considering the complex background of the rails and the presence of a large amount of non-defect interference information under the influence of the special lighting environment of the rails, a lightweight S-YOLOv8 student model is constructed and distilled using the T-YOLOv8 teacher model for training. This allows the optimized lightweight S-YOLOv8 student model to be used for real-time detection of rail surface defects, resulting in the following advantages: significantly improved detection performance across all categories, especially for difficult-to-detect defect categories, with the improved model achieving higher precision and recall. The improved model also achieves more consistent performance across different defect categories, demonstrating greater robustness and less susceptibility to the influence of individual difficult-to-detect categories. The model distillation method enables the lightweight S-YOLOv8 student model to inherit the advantages of the T-YOLOv8 teacher model, significantly improving detection efficiency. BRIEF DESCRIPTION OF THE DRAWINGS

[0033] The accompanying drawings illustrate exemplary embodiments and constitute a part of the specification. Together with the description of the specification, they serve to explain exemplary implementation of the embodiments. The illustrated embodiments are for illustrative purposes only and do not limit the scope of the claims. Throughout the drawings, the same reference numerals designate similar, but not necessarily identical, elements.

[0034] Figure 1 A schematic flow chart of a lightweight real-time detection method for track surface defects in the present invention;

[0035] Figure 2 A schematic diagram of the model structure to improve the original YOLOv8 model;

[0036] Figure 3 A schematic diagram of the model structure of the improved T-YOLOv8 teacher model of the present invention;

[0037] Figure 4 A schematic diagram of the model structure of the improved lightweight S-YOLOv8 student model of the present invention;

[0038] Figure 5 This is the PR curve of the original YOLOv8 model;

[0039] Figure 6 This is the PR curve of the improved T-YOLOv8 teacher model of the present invention;

[0040] Figure 7 This is the PR curve of the improved lightweight S-YOLOv8 student model of the present invention;

[0041] Figure 8 This is a detection result diagram of a lightweight real-time detection method for track surface defects in the present invention;

[0042] Figure 9 The figure is a structural schematic diagram of a lightweight real-time detection device for track surface defects in the present invention. DETAILED DESCRIPTION

[0043] In order to make the purpose, technical solutions and advantages of the present invention more clearly understood, the present invention is further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative work are within the scope of protection of the present invention.

[0044] It should be noted that the descriptions of "first", "second", etc. in the embodiments of the present invention are for descriptive purposes only and should not be understood as indicating or implying their relative importance or implicitly indicating the number of the indicated technical features. Therefore, the features defined as "first" or "second" may explicitly or implicitly include at least one of such features. In addition, the technical solutions between the various embodiments can be combined with each other, but this must be based on the fact that they can be implemented by ordinary technicians in this field. When the combination of technical solutions is contradictory or cannot be implemented, it should be deemed that such a combination of technical solutions does not exist and is not within the scope of protection required by the present invention.

[0045] In the description of the present invention, it should be understood that the numerical labels before the steps do not identify the order in which the steps are executed, but are only used to facilitate the description of the present invention and to distinguish each step, and therefore should not be understood as a limitation of the present invention.

[0046] First, an explanation of the terms involved in the present invention is provided:

[0047] YOLOv8 Network: YOLOv8 is the latest version of the YOLO family of object detection algorithms, focusing on improving detection speed and accuracy while introducing advanced network architecture and multi-task learning capabilities. It is suitable for real-time applications and offers strong flexibility to meet diverse detection needs.

[0048] HGNet module: It is mainly used to process features efficiently, improve model performance through multi-scale feature fusion, and improve efficiency by reducing the number of parameters and computational complexity, thereby achieving lightweight goals.

[0049] SPPF module: In the YOLO series and other target detection models, the SPPF module is often used to extract multi-scale features, thereby improving the performance and efficiency of the model by reducing the size of the feature map and retaining important information.

[0050] CBAM module: A dual attention mechanism module that adjusts feature maps in the channel dimension and spatial dimension respectively to highlight important information and suppress irrelevant information.

[0051] C2f module: It is a feature extraction and fusion module in deep learning, which contains the mechanism of channel attention and feature fusion.

[0052] CBAM-C2f module: The combination of CBAM and C2f further enhances the feature fusion capability.

[0053] Conv module: convolution module, used to extract features.

[0054] New-IAT module: A new image enhancement model with expanded convolution.

[0055] Dilated convolution layer: A convolution module that increases the receptive field by introducing holes (i.e., gaps) in the convolution kernel without increasing the computational burden or the number of parameters.

[0056] Detection head: The main component of the detection network, which is used to extract features and perform object classification and bounding box regression tasks.

[0057] PEM module: A branch of the New-IAT model, mainly used to gradually encode and process input image features to improve the model's performance.

[0058] GPM module: A branch of the New-IAT model that effectively captures and utilizes global features and contextual information to enhance the ability to understand and extract image features.

[0059] HGStem module: The HGStem (Hourglass Stem) module is a component of HGNet and is mainly used to extract preliminary features.

[0060] Light_HGBlock module: It is a lightweight version of the Hourglass block in HGNet, designed to reduce computational complexity and the number of parameters.

[0061] DWConv module: Depthwise Separable Convolution, which is a lightweight convolution operation.

[0062] Cross-attention mechanism: Cross-attention to different node features to enhance the feature expression ability of the detection target.

[0063] Multi-head attention mechanism: Expands the attention mechanism into multiple parallel "heads" to enhance the model's ability to capture and express information.

[0064] Secondly, in order to facilitate those skilled in the art to understand the technical solutions provided in the embodiments of the present invention, the technical solutions in the embodiments of the present invention will be described below with reference to the accompanying drawings:

[0065] Example 1

[0066] See Figure 1 In an embodiment of the present invention, a lightweight real-time detection method for track surface defects includes:

[0067] S1. Obtain track images to form a dataset, which is divided into a training set, a test set, and a validation set.

[0068] Specifically, images of railway tracks are collected through cameras and other image acquisition devices. After data cleaning and data labeling, the labeled track images are divided into training set, test set, and validation set according to preset ratios.

[0069] Data cleaning ensures that the acquired image dataset covers common track defect types, such as missing and broken fasteners, and trackbed debris. After detailed annotation of the cleaned image data, the training, test, and validation sets are divided in a ratio of 7:2:1.

[0070] Data annotation can be performed using annotation tools. Specifically, use the annotation tool LabelImg to perform the following image annotation process: create labels in LabelImg, and then convert the annotation information into the txt file format required for model training, where each image corresponds to a .txt file, which contains the category label of the target and the coordinate information of its bounding box.

[0071] S2. Use the dataset to pre-train the T-YOLOv8 teacher model.

[0072] Specifically, the T-YOLOv8 teacher model is based on the original YOLOv8 model. The NEW-IAT module is introduced into the backbone network, and the neck network is replaced with the CBAM-C2f module, which includes a dilated convolutional layer, a channel attention module, and a spatial attention module. Its structure and feature extraction process are as follows: the image input first passes through a 1x1 dilated convolutional layer to generate a feature map. This feature map is then split into two parts using a two-branch structure: the first part is passed directly, while the second part passes through several 3x3 dilated convolutional layers for feature extraction. These two parts are then concatenated along the channel dimension to generate a new feature map. This concatenation then enters the CBAM module, where it is first weighted along the channel dimension by the channel attention module and then weighted along the spatial dimension by the spatial attention module to generate an enhanced feature map. Finally, this enhanced feature map passes through a 1x1 convolutional layer to output the final feature map.

[0073] Furthermore, the NEW-IAT module includes a first branch and a second branch, the first branch includes an expanded convolution layer and a PEM module, and the second branch includes a convolution layer and an improved GPM module, wherein the PEM module is used to enhance local feature details, and the improved GPM module is used to generate global adjustment parameters.

[0074] contrast Figure 2 and Figure 3 From the model structure shown, we can see that the improvements to building the T-YOLOv8 teacher model are as follows:

[0075] 1. Improved Backbone Network: The NEW-IAT modules in the backbone network include but are not limited to: the first branch and the second branch. The first branch includes but is not limited to: the dilated convolutional layer and the PEM module. The second branch includes but is not limited to: the dilated convolutional layer and the improved GPM module. The improved GPM module uses a multi-head attention mechanism to replace the original cross-attention mechanism. The PEM module is responsible for enhancing local feature details, while the improved GPM module is responsible for generating global adjustment parameters. Together, these two modules provide more refined control for image enhancement.

[0076] 2. Improve the neck network: The CBAM-C2f modules of the neck network include but are not limited to: dilated convolutional layers, channel attention modules, and spatial attention modules. Adding these modules can speed up training and reduce the interference of redundant information.

[0077] 3. Fine-tune its parameters to optimize its performance on the railway track dataset.

[0078] 4. Use the aforementioned track dataset for pre-training to ensure that the model can effectively learn and adapt to various defect characteristics in the railway environment.

[0079] Further optionally, the specific steps of pre-training the T-YOLOv8 teacher model include:

[0080] S2.1. Use the NEW-IAT module to perform dark light enhancement or exposure correction on the image.

[0081] To achieve lightweightness, the original IAT model simplified its feature extraction process by using convolution and attention mechanisms, resulting in insufficient feature expression and low accuracy. The local branch of the New-IAT module retains its structure, replacing the original convolution with dilated convolution. This structure performs initial feature extraction, enabling a better fusion of local and global features. In the global branch, dilated convolutions replace the original convolutions. Furthermore, a multi-head attention mechanism replaces the cross-attention mechanism to capture global information, allowing the model to simultaneously focus on image features across multiple feature subspaces, improving accuracy. After processing the image with the New-IAT, its brightness returns to normal, making features easier to extract.

[0082] S2.2. Use the CBAM-C2f module to perform feature extraction on the image processed in S2.1.

[0083] A dilated convolutional layer is established in the CBAM-C2f module. A 3×3 convolution is performed within the dilated convolutional layer to reduce the dimensionality of the original feature map, resulting in a reduced-dimensional feature map. Dilated convolution is then performed on the reduced-dimensional feature map to capture richer contextual information. The feature map is then further processed using the Channel Attention Mechanism (CAM) and the Spatial Attention Mechanism (SAM). First, CAM extracts inter-channel correlations through global average pooling and global max pooling. This is then processed through a shared multi-layer perceptron (MLP) to generate channel weights that enhance important features. Second, SAM, while preserving inter-channel feature information, emphasizes location information by performing pixel-by-pixel attention calculations on the feature map. The resulting feature map contains more comprehensive and accurate spatial and channel information, helping to improve the model's object detection accuracy.

[0084] S2.3. Input the extracted image features into the prediction head in the T-YOLOv8 teacher model to generate the final prediction results of object detection.

[0085] The features processed by the New-IAT and CBAM-C2f modules are fed into subsequent layers of YOLOv8, including the prediction head, to generate the final prediction results for object detection. By utilizing these improved modules during training, the T-YOLOv8 teacher model maintains efficient and accurate detection performance even in environments with drastic lighting changes, and provides high-quality distillation guidance for the lightweight S-YOLOv8 student model. The prediction head is also called the detection head.

[0086] S3. Use the trained T-YOLOv8 teacher model to perform distillation training on the lightweight S-YOLOv8 student model, and optimize and verify the trained lightweight S-YOLOv8 student model.

[0087] Specifically, contrast Figure 2 and Figure 4 The lightweight S-YOLOv8 student model is based on the original YOLOv8 model, retains the SPPF layer, and replaces the rest of the backbone network with an improved lightweight new backbone network. The lightweight new backbone network includes three modules: HGStem, Light_HGBlock and DWConv.

[0088] Furthermore, the T-YOLOv8 teacher model is used to perform distillation training on the lightweight S-YOLOv8 student model, and model distillation technology is used to further optimize the lightweight S-YOLOv8 student model to improve its performance and efficiency in railway track defect detection tasks.

[0089] Prepare the T-YOLOv8 teacher model and the lightweight S-YOLOv8 student model: Select the pre-trained T-YOLOv8 teacher model, which has been optimized and adjusted to have high detection accuracy and generalization ability; select the lightweight S-YOLOv8 student model. The lightweight S-YOLOv8 student model is designed to maintain detection accuracy while reducing the model's computational complexity and memory consumption.

[0090] Specifically, the specific steps of distillation training include:

[0091] S3.1. Based on the classification loss, positioning loss and distillation loss, a distillation loss function is designed.

[0092] Distillation loss refers to the loss between the outputs of the T-YOLOv8 teacher model and the lightweight S-YOLOv8 student model. Classification loss calculates the difference between the class distribution of each detection box and the true label. Localization loss measures the difference in position between the bounding box predicted by the model and the true bounding box. By combining these losses into an overall loss function and using a weighted approach to balance the impact of each loss, we can ensure that the small model fully learns the knowledge of the large model while training it, thereby improving the performance and generalization ability of the small model.

[0093] S3.2. During the distillation training process, the track image is used as input, and iterative training is performed by adjusting the weights in the distillation loss function until the number of iterations meets the threshold.

[0094] Specifically, the above distillation training process includes: 1) configuring training hyperparameters, that is, parameters that need to be set in advance during the training process, with the input image resolution being 460, the batch size being 8, and the threshold for the number of iterations being 300 rounds; 2) knowledge extraction: performing forward inference through the T-YOLOv8 teacher model to obtain its output results on the training set, namely, feature maps and category probabilities; 3) soft label generation: using the feature maps and category probabilities output by the T-YOLOv8 teacher model as soft labels to guide the training of the lightweight S-YOLOv8 student model; 4) lightweight S-YOLOv8 student model training: training the lightweight S-YOLOv8 student model using soft labels. During the training process, training indicators, including loss values ​​and accuracy indicators, are regularly output, and verified on the validation set to obtain the model's performance indicators.

[0095] Specifically, after completing distillation training, the optimization and verification of the lightweight S-YOLOv8 student model include:

[0096] S3.3. Model optimization: Use optimization strategies to optimize the lightweight S-YOLOv8 student model to improve the performance of the lightweight S-YOLOv8 student model.

[0097] Among them, the optimization strategy is to combine the conventional target detection loss function and adjust the weight of the lightweight S-YOLOv8 student model so that it can better fit the output of the T-YOLOv8 teacher model during the learning process.

[0098] S3.4. Effect evaluation and verification: Use an independent test set to evaluate the performance of the lightweight S-YOLOv8 student model trained after distillation.

[0099] Specifically, by comparing indicators such as the accuracy, recall rate and F1 score of the model on the test set, the feasibility and effectiveness of the model in practical applications are verified.

[0100] Further optionally, in some embodiments, the specific operations of the model optimization in S3.3 above are as follows:

[0101] Batch load data, use Binary Cross-Entropy Loss function and The combination of losses is used to calculate the classification error and optimize the classification and localization accuracy.

[0102] The aforementioned binary cross entropy loss is calculated as follows:

[0103] L = 1 N ∑ i L i = − 1 N ∑ i [ y i ∗ l o g ( p ( y i ) ) + ( 1 − y i ) ∗ l o g ( 1 − p ( y i ) ) ]

[0104] Among them, L represents the binary cross entropy loss, N is the total number of samples, and i is the sample number. is the true label of the i-th sample, p( ) is the probability that the output belongs to the label.

[0105] The aforementioned CIoU loss function is as follows:

[0106]

[0107] in, represents CIoU loss, IoU represents the confidence threshold, c is the diagonal distance of the minimum bounding rectangle, Represents the predicted box b and the true box The center point distance between them, v is the correction factor considering the shape and direction of the target frame, is a hyperparameter.

[0108] The aforementioned DFL loss calculation formula is as follows:

[0109]

[0110] in, represents the DFL loss, and They are respectively the "predicted value" and "near predicted value" output by the network, 、 、 They are the “actual value”, “label integral value” and “adjacent label integral value” of the label respectively.

[0111] The specific process of converting labels into DFL format is as follows: convert the label value from xywh to ltrb; calculate the four values ​​of the label ltrb and convert them into The required integration form (i.e., label value + adjacent label values).

[0112] The specific conversion process is: y = the distance from the center to a certain edge / the current downsampling multiple.

[0113] In actual training, DFL is combined with CIoU loss. In the first step, DFL calculates the loss between the "bounding box distribution probability" and the "label distribution probability," thereby optimizing each edge. In the second step, the "bounding box distribution probability" is converted back to the predicted box, and the CIoU loss is used to calculate the loss between the predicted box and the "ground truth box" of the label, thereby optimizing the predicted box as a whole.

[0114] Furthermore, the distillation loss is calculated as follows:

[0115]

[0116]

[0117] in is the Kullback-Leibler divergence (KL divergence), q is the output of the lightweight S-YOLOv8 student model, and p is the output of the T-YOLOv8 teacher model. α, β, and γ are hyperparameters that control the impact of various types of losses on the total loss. L represents the binary cross entropy loss. represents the CIoU loss, Represents distillation loss.

[0118] Furthermore, the specific calculation method of the above indicators is as follows: the positive examples are correctly classified as positive examples, expressed as TP, the positive examples are incorrectly classified as negative examples, expressed as FN, the negative examples are correctly classified as negative examples, expressed as TN, and the negative examples are incorrectly classified as positive examples, expressed as FP;

[0119] The specific calculation is as follows:

[0120] Accuracy:

[0121] Recall:

[0122] Accuracy:

[0123] Average Precision (mAP): This is used to calculate the average performance of the model across the entire dataset. mAP is obtained by calculating the average precision (AP) for each class and then averaging the AP values ​​across all classes. AP is calculated using the Precision-Recall curve (PR curve), which plots precision and recall at different confidence thresholds. The calculation formula is as follows:

[0124] A P = ∑ k = 0 k = n − 1 [ R e c a l l s ( k ) − R e c a l l s ( k + 1 ) * P r e c i s i o n s ( k ) ]

[0125]

[0126] In the formula, k=0 to k=n-1 represent the number of results retrieved for all indices k and n from 0 to n-1. Recall(k) represents the recall rate for the element with index k in the current sort order, and Recall(k+1) represents the recall rate for the next index in the current sort order. Recall is the ratio of correctly identified positive examples to all actual positive examples. Precision(k) represents the precision for the element with index k in the current sort order. Precision is the ratio of correctly identified positive examples to all positive examples.

[0127] S4. Use the optimized lightweight S-YOLOv8 student model to detect track surface defects in real time to obtain and display track defect information.

[0128] Specifically, track surface defects may be cracks, wear, track bed debris, etc.; display refers to displaying the detected track defect information through an interactive web interface.

[0129] Specifically, the system uses image acquisition equipment to inspect railway tracks in real time and displays detected track defect information on an interactive webpage. The system involves: first, installing a detection device along the railway line or on a train to capture real-time track images; second, the device periodically captures images of the railway track and transmits the image data to a backend processing system; finally, the backend system applies an optimized lightweight S-YOLOv8 student model to perform real-time inspection of the received images. The detection algorithm analyzes the track area in the image and identifies possible defects such as cracks, wear, and trackbed debris. Finally, the detected track defect information is displayed in real time on an interactive webpage, which users can access and view.

[0130] Furthermore, the web interface includes, but is not limited to, marking the location and type of detected defects, along with images and detailed descriptions of the defects. Users can monitor the condition of the railway tracks in real time through the web interface and take necessary maintenance and repair measures immediately when problems are discovered.

[0131] In this embodiment of the present invention, considering the complex background of the rails and the presence of a large amount of non-defect interference information under the influence of the special lighting environment of the rails, a lightweight S-YOLOv8 student model is constructed and distilled using the T-YOLOv8 teacher model. This allows the optimized lightweight S-YOLOv8 student model to be used for real-time detection of rail surface defects, achieving the following advantages: significantly improved detection performance across all categories, especially for difficult-to-detect defect categories, with the improved model achieving higher precision and recall. The improved model also achieves more consistent performance across different defect categories, demonstrating greater robustness and less susceptibility to the influence of individual difficult-to-detect categories. This model distillation method enables the lightweight S-YOLOv8 student model to inherit the advantages of the T-YOLOv8 teacher model, significantly improving detection efficiency.

[0132] Specifically, first, the New-IAT dark light enhancement module was introduced into the T-YOLOv8 teacher model, enabling the model to better cope with illumination changes and improve recognition capabilities. At the same time, the CBAM-C2f module was introduced to address issues such as feature omissions and slow detection speeds. The dual attention mechanism in the module performs dual refinement on input features, enabling the model to simultaneously focus on meaningful channels and spatial locations. The adaptive recalibration mechanism helps the model better focus on key features, improving recognition performance and, in turn, speeding up the computational network. Secondly, HGNet was used in the lightweight S-YOLOv8 student model to replace the original complex and large convolution operations, making it more lightweight, facilitating model deployment and improving real-time detection efficiency. Ultimately, these model improvements resolved the issues of incorrectly detecting targets, slow detection speed, low accuracy, and difficult deployment in the traditional YOLOv8 model for rail surface defect detection scenarios.

[0133] It should also be noted that this invention employs a lightweight S-YOLOv8 student model and, through model distillation training, enhances the model's robustness under complex lighting conditions, making it applicable to a variety of practical working environments. While maintaining detection accuracy, it significantly improves the system's operating speed and responsiveness. Furthermore, this improves the accuracy and efficiency of railway track defect detection, making this solution even more effective in track object detection. It has broad application prospects in areas such as railway maintenance and safety management, enabling timely detection of railway track defects, ensuring the safety and reliability of railway systems, and effectively reducing labor and maintenance costs.

[0134] Based on the above embodiment method, the test was conducted on the test set, and the results are as follows Figure 5 、 6 As shown in Figure 7, the PR curve is used as the railway track detection evaluation index. Among them, mAP@0.5 means that the IoU (Intersection over Union) threshold used when calculating AP is 0.5. The PR curve is the precision (Precision) vs. recall rate (Recall) curve, with the recall rate as the horizontal axis and the precision as the vertical axis. According to Figure 5 、 6 ,7 before and after comparison can be found:

[0135] exist Figure 5 In the analysis, the Precision-Recall curve of the original YOLOv8 model is as follows:

[0136] 1) The mAP@0.5 for all categories is 0.932, indicating relatively high overall performance.

[0137] 2) The curves of different categories show some differences. For example, the performance of the trackbed debris category is significantly lower than that of other categories (both precision and recall are low, with an mAP of 0.779).

[0138] exist Figure 6 In the figure, the precision-recall curve of the improved T-YOLOv8 teacher model is analyzed as follows:

[0139] 1) mAP@0.5 for all categories increased to 0.991;

[0140] 2) The precision and recall rates of all categories have been improved, especially for the railbed debris category, where mAP@0.5 has increased to 0.982, a significant improvement.

[0141] 3) The overall curve is closer to the upper left corner, indicating that the overall performance of the model is stronger, with lower false positive and missed detection rates;

[0142] exist Figure 7 In the figure, the precision-recall curve of the lightweight S-YOLOv8 student model after distillation training is analyzed as follows:

[0143] 1) The mAP@0.5 for all categories is 0.992, which is close to the T-YOLOv8 teacher model.

[0144] 2) The performance of each category is similar to that of the T-YOLOv8 teacher model. The mAP for the track debris category is 0.978, which is slightly lower than the T-YOLOv8 teacher model but still significantly higher than the original YOLOv8 model.

[0145] 3) The precision and recall curves are similar to those of the T-YOLOv8 teacher model, indicating that the distillation effect is good and the lightweight S-YOLOv8 student model inherits most of the advantages of the T-YOLOv8 teacher model.

[0146] Based on the above Figure 5-7 The overall comparison is as follows:

[0147] 1) Overall Performance: The improved T-YOLOv8 teacher model and the distilled lightweight S-YOLOv8 student model both show significant improvements over the original models, with the overall mAP@0.5 increasing from 0.932 to 0.991 and 0.992.

[0148] 2) Single Category Performance: Categories that performed poorly in the original model showed significant improvements in both the T-YOLOv8 teacher model and the lightweight S-YOLOv8 student model. For example, the mAP for the railbed debris category increased from 0.779 to 0.982 (T-YOLOv8 teacher model) and 0.978 (lightweight S-YOLOv8 student model).

[0149] 3) Precision-Recall Curves: The precision and recall curves of the improved models (teacher and student) for each category are closer to the ideal state (upper left corner), indicating that both the false positive rate and the missed detection rate have decreased.

[0150] The lightweight S-YOLOv8 student model is able to approach or reach the performance level of the T-YOLOv8 teacher model while maintaining a small model size.

[0151] Overall, the improved model demonstrates significant improvements in detection performance across all categories, particularly for difficult-to-detect defect categories, achieving higher precision and recall. The improved model also demonstrates more consistent performance across different defect categories, demonstrating greater robustness and less vulnerability to individual hard-to-detect categories. The model distillation method enables the lightweight S-YOLOv8 student model to inherit the strengths of the T-YOLOv8 teacher model, significantly improving detection efficiency.

[0152] After this, the lightweight S-YOLOv8 student model is easier to deploy and the detection is faster. When the accuracy is close, the lightweight S-YOLOv8 student model is finally used for detection.

[0153] like Figure 8 The figure shows an example of real-time detection results based on the improved model of the present invention. As can be seen from the figure, the model can accurately detect the target and defect type, and the defect category is directly displayed in the result graph.

[0154] Example 2

[0155] like Figure 9 As shown, a lightweight real-time detection device for track surface defects in an embodiment of the present invention includes: an image acquisition module, a data processing module, a communication module, a power management module and a display module, wherein the data processing module is used to execute all the operating steps described in the above method embodiment to realize real-time detection of track surface defects.

[0156] Specifically, the image acquisition module may include a high-definition camera and a lighting system. The high-definition camera is used to capture track images and is installed at the bottom of the patrol vehicle to ensure that the field of view covers the entire track; the lighting system is used to provide fill light when light is insufficient to ensure image quality.

[0157] The data processing module includes an embedded computing unit and a storage unit. The embedded computing unit is used for image processing and model inference, and a high-performance embedded processor is selected to ensure real-time processing capabilities. The storage unit is used to store collected image data and detection results, and supports large-capacity storage.

[0158] The communication module includes a wireless communication unit and a GPS module. The wireless communication unit supports 4G / 5G or Wi-Fi communication and is used to transmit detection data to the monitoring center in real time; the GPS module is used to record the location of the detection device to help locate the defect.

[0159] The power management module includes a battery pack and a power management system. The battery pack provides power for the entire detection device and uses large-capacity batteries to ensure long-term stable operation. The power management system is responsible for battery charge and discharge management to ensure the reliability and safety of the equipment.

[0160] The display module mainly includes: a display for displaying the detection process and results.

[0161] Example 3

[0162] An embodiment of the present invention also provides a computer-readable storage medium having a computer program stored thereon, wherein when the computer program is executed by a processor, the steps of a lightweight real-time detection method for track surface defects in the embodiment are implemented.

[0163] In this embodiment, computer-readable storage media include flash memory, hard disks, multimedia cards, card-type memories (e.g., SD or DX memories), random access memory (RAM), static random access memory (SRAM), read-only memory (ROM), electrically erasable programmable read-only memory (EEPROM), programmable read-only memory (PROM), magnetic storage, magnetic disks, optical disks, and the like. In some embodiments, the computer-readable storage medium may be an internal storage unit of a computer device, such as the computer device's hard disk or memory. In other embodiments, the computer-readable storage medium may also be an external storage device of the computer device, such as a plug-in hard disk, a Smart Media Card (SMC), a Secure Digital (SD) card, a flash memory card, and the like. Of course, the computer-readable storage medium may also include both the internal storage unit and external storage devices of the computer device. In this embodiment, the computer-readable storage medium is typically used to store the operating system and various application software installed on the computer device, such as the program code for a lightweight real-time rail surface defect detection method described in the embodiment. Furthermore, the computer-readable storage medium may also be used to temporarily store various types of data that has been or will be output.

[0164] Example 4

[0165] An embodiment of the present invention further provides a computer program product, including a computer program, which implements the method in the above embodiment when executed by a processor.

[0166] Obviously, those skilled in the art should understand that the various modules or steps of the above-mentioned embodiments of the present invention can be implemented using general-purpose computer devices, they can be concentrated on a single computer device, or distributed on a network composed of multiple computer devices. Alternatively, they can be implemented using program codes executable by the computer device, so that they can be stored in a storage device and executed by the computer device. In some cases, the steps shown or described can be performed in a different order than herein, or they can be made into individual integrated circuit modules, or multiple modules or steps can be made into a single integrated circuit module for implementation. Thus, the embodiments of the present invention are not limited to any specific combination of hardware and software.

[0167] It should be noted that the above are only preferred embodiments of the present invention and do not limit the scope of patent protection of the present invention. Any equivalent structure or equivalent process transformation made using the contents of the description and drawings of the present invention, or directly or indirectly applied in other related technical fields, are also included in the scope of patent protection of the present invention.

Claims

1. A lightweight real-time detection method for track surface defects, characterized in that: include: S1. Acquire track images to form a dataset, which is divided into a training set, a test set, and a validation set; S2. Pre-train the T-YOLOv8 teacher model using the dataset; The T-YOLOv8 teacher model is based on the original YOLOv8 model. The NEW-IAT module is introduced into the backbone network and the C2f module in the neck network is replaced by the CBAM-C2f module. S3. Use the trained T-YOLOv8 teacher model to perform distillation training on the lightweight S-YOLOv8 student model, and optimize and verify the trained lightweight S-YOLOv8 student model; The lightweight S-YOLOv8 student model is based on the original YOLOv8 model, retaining the SPPF layer and replacing the original YOLOv8 backbone network with a lightweight new backbone network. The lightweight new backbone network includes the HGStem module, the Light_HGBlock module, and the DWConv module. S4. Use the optimized lightweight S-YOLOv8 student model to detect rail surface defects in real time to obtain and display rail defect information; The NEW-IAT module includes a first branch and a second branch. The first branch includes an expanded convolution layer and a PEM module. The second branch includes an expanded convolution layer and an improved GPM module. The improved GPM module uses a multi-head attention mechanism to replace the original cross attention mechanism. The PEM module is used to enhance local feature details, and the improved GPM module is used to generate global adjustment parameters. The CBAM-C2f module includes a dilated convolutional layer, a channel attention module, and a spatial attention module. Its structure and feature extraction process are as follows: the image input first passes through a 1x1 dilated convolutional layer to generate a feature map; then the feature map is divided into two parts through a two-branch structure: the first part is directly passed, and the second part passes through several 3x3 dilated convolutional layers for feature extraction. The two parts of the feature are then spliced ​​in the channel dimension to obtain a new feature map; the new feature map after splicing enters the CBAM module, first passes through the channel attention module to perform channel weighting, and then passes through the spatial attention module to perform spatial weighting to obtain an enhanced feature map. Finally, the enhanced feature map passes through a 1x1 convolutional layer to output the final feature map; In S2, the specific steps for pre-training the T-YOLOv8 teacher model include: S2.

1. Use the NEW-IAT module to perform dark light enhancement or exposure correction on the image. S2.2, use the CBAM-C2f module to perform feature extraction on the image processed in S2.1; S2.

3. Input the extracted image features into the prediction head of the T-YOLOv8 teacher model to generate the final prediction results of object detection; In S3, the specific steps of distillation training include: S3.

1. Design a distillation loss function based on classification loss, localization loss, and distillation loss. S3.

2. During the distillation training process, the track image is used as input, and iterative training is performed by adjusting the weights in the distillation loss function until the number of iterations meets the threshold.

2. A lightweight real-time detection method for track surface defects according to claim 1, characterized in that: After completing distillation training, the optimization and verification of the lightweight S-YOLOv8 student model include: S3.

3. Model Optimization: Use an optimization strategy to optimize the lightweight S-YOLOv8 student model to improve the performance of the lightweight S-YOLOv8 student model. The optimization strategy is to adjust the weights of the lightweight S-YOLOv8 student model in combination with the conventional object detection loss function. S3.

4. Effect evaluation and verification: Use an independent test set to evaluate the performance of the lightweight S-YOLOv8 student model trained after distillation.

3. A lightweight real-time detection method for track surface defects according to claim 1, characterized in that: Display is to show the detected track defect information through an interactive web interface; The web interface includes annotations of the detected defect location and type, defect images and detailed descriptions.

4. A lightweight real-time detection method for track surface defects according to claim 1, characterized in that: The obtained track images form a data set, which is divided into a training set, a test set, and a validation set, including: Railway track images are collected using cameras and other image acquisition devices. After data cleaning and labeling, the labeled track images are divided into training set, test set, and validation set according to preset ratios.

5. A lightweight real-time detection method for track surface defects according to claim 4, characterized in that: The LabelImg annotation tool is used for data annotation. The following image annotation process is performed using LabelImg: labels are created in LabelImg, and then the annotation information is converted into the txt file format required for model training. Each image corresponds to a .txt file, which contains the category label of the target and the coordinate information of its bounding box.

6. A lightweight real-time detection device for track surface defects, characterized in that: include: An image acquisition module, a data processing module, a communication module, a power management module and a display module, wherein the data processing module is used to execute a lightweight real-time detection method for track surface defects as described in any one of claims 1 to 5.

Images