Lightweight track surface defect real-time detection method and device

Abstract

The invention discloses a lightweight track surface defect real-time detection method and device, relates to the technical field of image data processing, and improves the detection precision and detection efficiency of track defect detection by optimizing a model structure and a training strategy. The method comprises the following steps: pre-training a T-YOLOv8 teacher model by using a data set of a track image; the T-YOLOv8 teacher model is obtained by taking an original YOLOv8 model as a basis, introducing an NEW-IAT module into a backbone network and replacing a neck network with a CBAM-C2f module; carrying out distillation training on the lightweight S-YOLOv8 student model by using the trained T-YOLOv8 teacher model, and carrying out optimization and verification on the trained lightweight S-YOLOv8 student model; the lightweight S-YOLOv8 student model is based on an original YOLOv8 model, an SPPF layer is reserved, the rest part of a backbone network is replaced by a lightweight new backbone network, and the lightweight new backbone network comprises three modules of HGStem, LightHGBlock and DWConv; and the optimized lightweight S-YOLOv8 student model is used to carry out real-time detection on 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 the long-term use, the rails on the track lines are affected by the dual influence of 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 time, they are very likely to lead to serious accidents such as train derailment and collision, causing huge casualties and property losses, and bringing huge challenges to the operational guarantee and further healthy development of rail transit.

[0003] With the development of computer vision and deep learning technology, automated track defect detection based on image and video data has become an effective solution to track defects. In the specific task of railway defect detection, since there are many types of defects on the railway surface and their shapes vary, how to effectively and accurately identify and classify these different types of defects is a difficult problem. In addition, the background of the railway track is complex and there is a large amount of non-defect interference information. How to accurately locate and identify tiny defects in this complex background is also a problem that needs to be solved. It is precisely because of the existence of the above difficulties that the existing detection technology still has problems of low detection accuracy and low detection efficiency. 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, obtaining orbital images to form a data set, which is divided into a training set, a test set, and a validation set;

[0007] S2, use the dataset to pre-train the T-YOLOv8 teacher model;

[0008] Among them, 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 an expanded 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] Among them, the lightweight S-YOLOv8 student model is based on the original YOLOv8 model, retains the SPPF layer and replaces the original backbone network of YOLOv8 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 passing it through 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 through several 3x3 dilated convolution layers. The two parts of the feature are then concatenated in the channel dimension to obtain a new feature map; the new feature map after concatenation enters the CBAM module, firstly, it is weighted in the channel dimension through the channel attention module, and then it is weighted in the spatial dimension through the spatial attention module to obtain an enhanced feature map, and 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 target 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, positioning 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 training is completed when the number of iterations meets the threshold.

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

[0023] S3.3, model optimization: use the optimization strategy to optimize the lightweight S-YOLOv8 student model to improve the performance of the lightweight S-YOLOv8 student model; wherein the optimization strategy is to adjust the weight of the lightweight S-YOLOv8 student model in combination with the conventional target 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 with 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 page interface; the web page interface includes marking the detected defect location and type, defect pictures and detailed descriptions.

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

[0027] The railway track images 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.

[0028] Optionally, in a possible implementation of the first aspect, a 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 a 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 first aspect and its implementation methods is implemented.

[0031] In the fourth aspect, another aspect of the present invention provides a computer program product, including 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 technical scheme adopted in the present invention can achieve the following technical effects: considering the special lighting environment of the track, the background of the rail is complex and there is a large amount of non-defect interference information, a lightweight S-YOLOv8 student model is constructed, and the student model is distilled and trained using the T-YOLOv8 teacher model, so that the optimized lightweight S-YOLOv8 student model is used to perform real-time detection of track surface defects to obtain the following advantages: the detection performance of all categories is significantly improved, especially for difficult-to-detect defect categories, the effect is particularly significant, and the improved model has higher precision and recall rate; the performance of the improved model on different defect categories is more consistent, indicating that it is more robust and not easily affected by 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 well, and the detection efficiency is significantly improved. BRIEF DESCRIPTION OF THE DRAWINGS

[0033] The accompanying drawings exemplarily illustrate the embodiments and constitute a part of the specification, and together with the text description of the specification, are used to explain the exemplary implementation of the embodiments. The embodiments shown are for illustrative purposes only and do not limit the scope of the claims. In all drawings, the same reference numerals refer to 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 It is the PR curve diagram of the improved T-YOLOv8 teacher model of the present invention;

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

[0041] Figure 8 A detection result diagram of a lightweight rail surface defect real-time detection method in the present invention;

[0042] Fig. 9 The figure is a schematic structural 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 scheme and advantages of the present invention more clearly understood, the present invention is further described in detail below in conjunction with 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 involving "first", "second", etc. in the embodiments of the present invention are only for descriptive purposes and cannot be understood as indicating or suggesting their relative importance or implicitly indicating the number of the indicated technical features. Therefore, the features defined as "first" and "second" may explicitly or implicitly include at least one of the features. In addition, the technical solutions between the various embodiments can be combined with each other, but they must be based on the ability of ordinary technicians in the field to implement them. 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 of executing the steps, but are only used to facilitate the description of the present invention and distinguish each step, and therefore cannot be understood as a limitation of the present invention.

[0046] First, the following terms are explained:

[0047] YOLOv8 network: YOLOv8 is the latest version of the YOLO series of target 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 application scenarios, has strong flexibility, and can meet different detection needs.

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

[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: The attention mechanism is expanded to 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 in conjunction with the accompanying drawings:

[0065] Embodiment 1

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

[0067] S1. Acquire orbital images to form a data set, 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 can ensure that the acquired image data set covers common types of track defects, such as missing fasteners, fractures, and trackbed debris. After detailed annotation of the cleaned image data, the training set, test set, and validation set 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. Pre-train the T-YOLOv8 teacher model using the dataset.

[0072] Specifically, the T-YOLOv8 teacher model is based on the original YOLOv8 model, with the NEW-IAT module introduced into the backbone network and the neck network replaced by the CBAM-C2f module. 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: the image input first passes through a 1x1 expanded convolution 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 expanded convolution 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, firstly, it is weighted in the channel dimension through the channel attention module, and then weighted in the spatial dimension through the spatial attention module to obtain an enhanced feature map, and finally, the enhanced feature map passes through a 1x1 convolution 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, it can be seen that the improvements in 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: dilated convolution layer, PEM module, and the second branch includes but is not limited to: dilated convolution layer, 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. The two modules provide more refined control for image enhancement;

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

[0077] 3. Fine-tune its parameters to optimize its effect 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] In order to achieve lightweight, the original IAT model simplified the convolution and attention mechanism in its feature extraction process, resulting in insufficient feature expression ability and low feature expression accuracy. The Local branch of the New-IAT module uses dilated convolution to replace the original convolution, and the other structures remain unchanged. Dilated convolution is used for preliminary feature extraction, and dilated convolution can better integrate local and global features. In the Global branch, dilated convolution is also used to replace the original convolution. At the same time, the multi-head attention mechanism is used to replace the cross attention mechanism to capture global information, allowing the model to focus on image features in multiple different feature subspaces in parallel, which can improve its accuracy. After the image is processed by New-IAT, its brightness tends to be normal and features are easier to extract.

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

[0083] In the CBAM-C2f module, a dilated convolution layer is established. A 3×3 convolution is used in the dilated convolution layer to reduce the dimension of the original feature map to obtain a reduced-dimensional feature map. The dilated convolution is performed on the reduced-dimensional feature map to capture richer contextual information. Then, the feature map is further processed by the Channel Attention Mechanism (CAM) and the Spatial Attention Mechanism (SAM). First, CAM extracts the correlation between channels through global average pooling and global maximum pooling, and processes it through a shared multi-layer perceptron (MLP) to finally generate channel weights to enhance important features. Secondly, SAM emphasizes the focus on position information by performing pixel-by-pixel attention calculation on the feature map while maintaining the feature information between channels. The final output feature map contains more comprehensive and accurate spatial and channel information, which helps to improve the accuracy of the model in object detection.

[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 target detection.

[0085] The features processed by the New-IAT module and the CBAM-C2f module are input into the subsequent layers of YOLOv8, including the prediction head, to generate the final prediction results of target detection. By using these improved modules during the training process, the T-YOLOv8 teacher model can still maintain efficient and accurate detection performance in environments with drastic lighting changes, and provide 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 the model distillation technology is used to further optimize the lightweight S-YOLOv8 student model to improve its performance and efficiency in the railway track defect detection task.

[0089] Prepare T-YOLOv8 teacher model and 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 computational complexity and memory consumption of the model.

[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 output of the T-YOLOv8 teacher model and the lightweight S-YOLOv8 student model. Classification loss calculates the difference between the category distribution of each detection box and the true label. Localization loss measures the position difference between the bounding box predicted by the model and the true bounding box. By combining them into a total loss function and balancing the impact of each loss in a weighted manner, it is possible to ensure that the small model can fully learn 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 training is completed when the number of iterations meets the threshold.

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

[0095] Specifically, after completing the 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 weights 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 with 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 implementations, the specific operations of the model optimization in S3.3 above are as follows:

[0101] Batch load data, use binary cross-entropy loss function (Binary Cross-Entropy Loss) 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 calculation formula is 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 box, 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 "prediction value" and "near prediction 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 to DFL format is as follows: convert the label value from xywh to ltrb; calculate the four values ​​of the label's ltrb and convert them into The required integral form (i.e., label value + neighboring 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 used in combination with CIoU loss. In the first part, DFL is used to calculate the loss of the "bounding box distribution probability" and the "distribution probability" of the label, thereby optimizing each edge. In the second part, the "bounding box distribution probability" is restored to the predicted box, and the loss of the predicted box and the "actual box" of the label is calculated through CIoU loss, 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: correctly classify the positive example as a positive example, expressed as TP, misclassify the positive example as a negative example, expressed as FN, correctly classify the negative example as a negative example, expressed as TN, and misclassify the negative example as a positive example, expressed as FP;

[0119] The specific calculation is as follows:

[0120] Accuracy:

[0121] Recall:

[0122] Accuracy:

[0123] Average Precision (mAP): It is used to calculate the average performance of the model on the entire dataset. mAP is obtained by calculating the average precision (AP) for each class and then calculating the average AP value of all classes. The AP is calculated based on the Precision-Recall curve (PR curve), which calculates the precision and recall rate based on 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 represents the number of results retrieved by all indexes k and n from 0 to n-1, Recall(k) represents the recall rate corresponding to the element with index k in the current sort, and Recall(k+1) represents the recall rate of the next index in the current sort. Recall is the ratio of correctly identified positive examples to all actual positive examples. Precision(k) represents the precision corresponding to the element with index k in the current sort, and precision is the ratio of correctly identified positive examples to the total number of positive examples identified.

[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, the track surface defects may be cracks, wear, track bed debris, etc.; display refers to displaying the detected track defect information through an interactive web page interface.

[0129] Specifically, the image acquisition equipment is used to detect railway tracks in real time, and the detected track defect information is displayed through an interactive web page, including: first, the detection device is installed along the railway or on the train to capture real-time track images; second, the equipment regularly captures images of railway tracks and transmits the image data to the back-end processing system; finally, in the back-end system, the optimized lightweight S-YOLOv8 student model is applied to perform real-time detection of the received images. The detection algorithm analyzes the track area in the image and identifies possible defects such as cracks, wear, and roadbed debris. Finally, the detected track defect information is displayed in real time on an interactive web page, and users can access and view it through the web page.

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

[0131] In the embodiment of the present invention, considering the special lighting environment of the track, the background of the rail is complex and there is a large amount of non-defect interference information, a lightweight S-YOLOv8 student model is constructed, and the lightweight S-YOLOv8 student model is distilled and trained using the T-YOLOv8 teacher model, so that the optimized lightweight S-YOLOv8 student model is used to perform real-time detection of track surface defects to obtain the following advantages: significantly improve the detection performance of all categories, especially for difficult-to-detect defect categories, the effect is particularly significant, and the improved model has higher precision and recall rate; the improved model has more consistent performance on different defect categories, indicating that it is more robust and not easily affected by 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 well, and the detection efficiency is significantly improved.

[0132] Specifically, firstly, the New-IAT dark light enhancement module is introduced into the T-YOLOv8 teacher model, so that the model can better cope with the problem of illumination changes and improve the recognition ability; at the same time, the CBAM-C2f module is introduced to address the problems of feature omission and slow detection speed. The dual attention mechanism in the module performs dual refinement on the input features, so that the model can focus on meaningful channels and spatial positions at the same time. The adaptive recalibration mechanism helps the model to better focus on key features, improve recognition performance, and thus speed up the operation speed of the computing network. Secondly, HGNet is used in the lightweight S-YOLOv8 student model to replace the original complex and huge convolution operation, making it more lightweight, which is helpful for model deployment and improving real-time detection efficiency. Finally, through the above model improvements, the problems of wrong detection targets, slow detection speed, low accuracy, and difficult deployment of traditional YOLOv8 in the track surface defect detection scenario are solved.

[0133] It should also be noted that the present invention adopts the construction of a lightweight S-YOLOv8 student model, and through model distillation training, the robustness of the model under complex lighting conditions is enhanced, making it applicable to a variety of actual working environments, while maintaining the detection accuracy, significantly improving the system's operating speed and responsiveness. In turn, the accuracy and efficiency of railway track defect detection are improved, making this solution perform better in track target detection, and it can have broad application prospects in the fields of railway maintenance and safety management, and can be used to detect railway track defects in a timely manner, ensuring the safety and reliability of the railway system, thereby effectively reducing labor and maintenance costs.

[0134] Based on the above embodiment method, the test is performed 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 a relatively high overall performance.

[0137] 2) The curves of different categories show some differences. For example, the performance of the track 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 of each category have been improved, especially for the rail debris category, where mAP@0.5 has increased to 0.982, which is a significant effect;

[0141] 3) The overall curve is closer to the upper left corner, indicating that the overall performance of the model is stronger, and the false positive rate and missed positive rate are reduced;

[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 of 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 are significantly improved compared to 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 were significantly improved in both the T-YOLOv8 teacher model and the lightweight S-YOLOv8 student model. For example, the mAP of the rail debris category was improved from 0.779 to 0.982 (T-YOLOv8 teacher model) and 0.978 (lightweight S-YOLOv8 student model).

[0149] 3) Precision-Recall Curve: The precision and recall curves of the improved models (teacher and student) in each category are closer to the ideal state (upper left corner), indicating that both the false positive rate and the missed positive 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 advantages of the improved model are: significantly improved detection performance for all categories, especially for difficult-to-detect defect categories, the improved model has higher precision and recall; the improved model has more consistent performance on different defect categories, indicating that it is more robust and less susceptible to 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, and the detection efficiency is significantly improved.

[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 the detection result of real-time detection based on the improved model of the present invention. It can be seen from the figure that the model can accurately detect the target and defect type, and its defect category is also directly displayed in the result figure.

[0154] Embodiment 2

[0155] like Fig. 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 operation 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 reasoning. A high-performance embedded processor is selected to ensure real-time processing capabilities. The storage unit is used to store the 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 and 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 selects 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] Embodiment 3

[0162] A computer-readable storage medium is also provided in an embodiment of the present invention, on which a computer program is stored, 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, the computer-readable storage medium includes flash memory, hard disk, multimedia card, card-type memory (for example, SD or DX memory, etc.), 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 memory, disk, optical disk, etc. In some embodiments, the computer-readable storage medium may be an internal storage unit of a computer device, such as a hard disk or memory of the computer device. In other embodiments, the computer-readable storage medium may also be an external storage device of a computer device, such as a plug-in hard disk, a smart memory card (Smart Media Card, SMC), a secure digital (Secure Digital, SD) card, a flash card (Flash Card), etc. equipped on the computer device. Of course, the computer-readable storage medium may also include both the internal storage unit of the computer device and its external storage device. In this embodiment, the computer-readable storage medium is generally used to store an operating system and various application software installed on the computer device, such as a program code of a lightweight real-time detection method for track surface defects in the embodiment, etc. In addition, the computer-readable storage medium may also be used to temporarily store various types of data that have been output or are to be output.

[0164] Embodiment 4

[0165] A computer program product is also provided in an embodiment of the present invention, including a computer program, and when the computer program is executed by a processor, the method in the above embodiment is implemented.

[0166] Obviously, those skilled in the art should understand that the modules or steps of the above-mentioned embodiments of the present invention can be implemented by general-purpose computer devices, they can be concentrated on a single computer device, or distributed on a network composed of multiple computer devices, and optionally, they can be implemented by executable program codes of computer devices, so that they can be stored in a storage device and executed by the computer device, and in some cases, the steps shown or described can be executed in a different order from that here, or they can be made into individual integrated circuit modules, or multiple modules or steps therein can be made into a single integrated circuit module for implementation. In this way, 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 the patent protection scope of the present invention is not limited thereto. Any equivalent structure or equivalent process transformation made using the contents of the present invention specification and drawings, or directly or indirectly applied in other related technical fields, are also included in the patent protection scope of the present invention.

Claims

1. A lightweight real-time detection method for track surface defects, characterized in that: include: S1, obtaining orbital images to form a data set, which is divided into a training set, a test set, and a validation set; S2, use the dataset to pre-train the T-YOLOv8 teacher model; Among them, 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. Among them, the lightweight S-YOLOv8 student model is based on the original YOLOv8 model, retains the SPPF layer and replaces the original backbone network of YOLOv8 with a lightweight new backbone network. The lightweight new backbone network includes the HGStem module, Light_HGBlock module and DWConv module; S4. Use the optimized lightweight S-YOLOv8 student model to detect track surface defects in real time to obtain and display track defect information.

2. A lightweight real-time detection method for track surface defects according to claim 1, characterized in that: 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: the image input first passes through a 1x1 expanded convolution 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 expanded convolution 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, firstly, it is weighted in the channel dimension through the channel attention module, and then it is weighted in the spatial dimension through 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.

3. A lightweight real-time detection method for track surface defects according to claim 1, characterized in that: 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 in the T-YOLOv8 teacher model to generate the final prediction results of target detection.

4. A lightweight real-time detection method for track surface defects according to claim 1, characterized in that: In S3, the specific steps of distillation training include: S3.

1. Design a distillation loss function based on classification loss, positioning 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 training is completed when the number of iterations meets the threshold.

5. A lightweight real-time detection method for track surface defects according to claim 1 or 4, characterized in that: After completing the distillation training, the optimization and verification of the lightweight S-YOLOv8 student model include: S3.3, model optimization: use the optimization strategy to optimize the lightweight S-YOLOv8 student model to improve the performance of the lightweight S-YOLOv8 student model; wherein the optimization strategy is to adjust the weight of the lightweight S-YOLOv8 student model in combination with the conventional target 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 with distillation.

6. 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.

7. A lightweight real-time detection method for track surface defects according to claim 1, characterized in that: The acquired track images form a data set, which is divided into a training set, a test set, and a validation set, including: The railway track images 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.

8. A lightweight real-time detection method for track surface defects according to claim 7, characterized in that: When annotating data, use the annotation tool LabelImg. Use 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. Each image corresponds to a .txt file, which contains the category label of the target and the coordinate information of its bounding box.

9. 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 8.

Images