A track defect detection method based on task alignment and dynamic category weighting

Abstract

The application discloses a track defect detection method based on task alignment and dynamic category weighting, and belongs to the technical field of track transportation operation and maintenance, and the steps are as follows: a high-definition camera device is carried by a drone to perform inspection on a railway or an urban track to obtain track image data; the collected images are preprocessed and divided into a training set, a verification set and a test set; in the training process of a YOLO target detection network, task alignment scoring is introduced in the positive sample allocation stage; a unified rare-size joint weighting classification loss is designed, a dynamic updated category weight is introduced, and the influences of category rarity and target size on learning difficulty are simultaneously described, so that the overall detection precision is improved; and the trained model is used to infer new track images collected by the drone, and detection results containing defect categories, positions and confidence are output.

Description

Technical Field

[0001] This invention belongs to the field of rail transit operation and maintenance technology, specifically relating to a track defect detection method based on task alignment and dynamic category weighting. Background Technology

[0002] With the rapid development of railways and urban rail transit, the health of track structures directly affects train operation safety and maintenance efficiency. Traditional inspection methods, such as manual inspection and track inspection vehicle sensor detection, suffer from problems such as high labor intensity, low efficiency, long inspection cycles, and limited coverage. In recent years, drone-mounted camera inspection methods have been able to quickly acquire visible light images of the track, providing a data foundation for vision-based intelligent inspection.

[0003] Existing deep learning-based object detection methods (such as the YOLO series) have made progress in track defect recognition, but they still have the following common shortcomings: First, positive sample allocation often relies on the geometric overlap (IoU) between the predicted box and the ground truth box or a fixed threshold, ignoring the semantic confidence (classification score) of the candidate boxes. This easily leads to the inclusion of candidates that are "geometrically close but semantically unreliable" into the positive samples, resulting in high training noise and limited improvement in mAP@0.5:0.95. Second, track defects generally exhibit the characteristics of "long tail + small target". Rare class samples are few and greatly affected by background texture interference. Common unified loss weights are difficult to balance overall accuracy and rare class recall. Third, structural modifications or complex global matching (such as optimal transmission) made to improve performance often bring significant engineering intrusion and additional computation, affecting the implementation on existing inspection platforms.

[0004] Therefore, there is an urgent need for a training method that can balance geometric matching and semantic confidence without changing the network structure and training hyperparameters, and can reweight the classification loss for long-tail categories. On the one hand, it can suppress noisy positive samples and improve the effective supervision of small targets and rare classes during the sample allocation stage. On the other hand, it can minimize engineering changes and have zero additional overhead during the inference stage, thus facilitating rapid deployment in UAV inspection scenarios. Summary of the Invention

[0005] To address the aforementioned problems in existing technologies, this invention proposes a track defect detection method based on task alignment and dynamic category weighting. The method is rationally designed, overcomes the shortcomings of existing technologies, and has good performance.

[0006] To achieve the above objectives, the present invention adopts the following approach:

[0007] A track defect detection method based on task alignment and dynamic category weighting includes the following steps:

[0008] S1: Use drones equipped with high-definition cameras to inspect railways or urban tracks and acquire track image data;

[0009] S2: Preprocess the acquired images and divide them into training set, validation set and test set;

[0010] S3: In the training process of the YOLO object detection network, task alignment scoring is introduced in the positive sample allocation stage;

[0011] S4: Design a unified scarcity-size joint weighted classification loss, introduce dynamically updated class weights, and simultaneously characterize the impact of class scarcity and target size on learning difficulty, thereby improving overall detection accuracy;

[0012] S5: Use the trained model to infer the new orbital images acquired by the UAV and output the detection results including defect category, location and confidence level.

[0013] Furthermore, in S2, the preprocessing includes image annotation, image normalization, size scaling, brightness enhancement, and data augmentation;

[0014] Rail defects include spalling, wheel burns, rail head crushing, and corrugated wear;

[0015] Each track defect target is labeled according to the YOLO text format with class, x_center, y_center, width, and height. class is the category, x_center is the center x coordinate, y_center is the center y coordinate, width is the width, and height is the height. All coordinates are normalized relative values.

[0016] Further, step S3 includes the following sub-steps:

[0017] S3.1: For a training image, denote all predicted candidate boxes as a set. All truth boxes are sets All categories are set as Each truth box corresponds to a category. ;

[0018] For each candidate box-truth box pair Calculate the location item and the classification item separately; the location item is the geometric overlap. The classification item is aligned to the truth box category, and the calculation formula is as follows:

[0019] ;

[0020] in, Candidate boxes In category The logit value, which is the output value of the YOLO object detection network; Candidate boxes Category confidence level for function;

[0021] S3.2: For each Define a comprehensive score, and record the score as follows: The expression is:

[0022] ;

[0023] in, These are weighting coefficients;

[0024] S3.3: For each truth box ,right Sort the candidate boxes in descending order of score and select the top-k candidate boxes with the highest scores as positive candidate samples. If the same candidate box matches multiple ground truth boxes, only the matching relationship with the ground truth box with the highest score is retained.

[0025] Further, step S4 includes the following sub-steps:

[0026] S4.1: For each category Calculate the frequency of occurrence and the average normalized area, respectively. The frequency of occurrence is the category within the most recent iterations. The dynamic smoothing statistic for the selected positive sample is expressed as:

[0027] ;

[0028] in, and The frequencies of occurrence before and after the update are respectively. This represents the current batch of positive samples. These are the weighted smoothing coefficients. ;

[0029] The average normalized area is the category in recent rounds. The dynamic smoothed average of the normalized area of ​​positive samples is expressed as:

[0030] ;

[0031] in, This represents the mean area of ​​the current batch of positive samples.

[0032] S4.2: Combine the scarcity and size terms using a linear weighting to form the class weights. Obtain the category weight function:

[0033] ;

[0034] in, For weighting coefficients, adjust and Impact on the final weight, , To prevent the occurrence of constants that are divided by zero, , The parameters are used to adjust the influence of the category scarcity term and the target size term on the weights, respectively. ;

[0035] S4.3: To After normalization and clipping, the expression is:

[0036] ;

[0037] ;

[0038] in, The normalized class weights, and These are the minimum and maximum values ​​of the category weights, respectively. For the final category weight;

[0039] S4.4: Will A weighted classification loss is constructed by applying the classification supervision term to each positive sample. :

[0040] ;

[0041] in, Positive sample The true category label, Positive samples output by the model The predicted probability of the category. The binary cross-entropy loss function is... Positive sample The true class weights, For the positive sample set, The number of positive samples.

[0042] The beneficial technical effects of this invention are as follows:

[0043] This invention proposes a unified scarcity-size joint weighted positive sample classification loss without modifying the YOLO backbone and detection head. It is combined with a TAL-like Top-K positive sample competitive retention and batch-based positive sample normalization training strategy: adaptive weighting for few samples / small-sized defects, suppressing false detections caused by complex backgrounds such as ballast and weeds, improving the separability of positive and negative samples and convergence stability, and significantly improving mAP@0.5 and mAP@0.5:0.95. This invention only modifies the loss and matching stages; training and deployment are fully compatible with the original YOLO, making it highly practical for engineering applications. Attached Figure Description

[0044] Figure 1 This is a flowchart illustrating the overall process of the method of the present invention.

[0045] Figure 2 This is a schematic diagram of the task alignment and allocation mechanism of the present invention.

[0046] Figure 3 This is a schematic diagram of the track defect detection results of the method of the present invention.

[0047] Figure 4 This is a statistical diagram illustrating the distribution of track defect samples.

[0048] Wherein, (a) is the category distribution; (b) is the distribution of the target bounding box shape and center; (c) is the spatial distribution of the target center (x, y); and (d) is the distribution of the target size (width, height).

[0049] Figure 5 Precision-confidence curves for different defect categories.

[0050] Figure 6 Recall-confidence curves for different defect categories.

[0051] Figure 7 F1 score-confidence curves for different defect categories. Detailed Implementation

[0052] The specific embodiments of the present invention will be further described below with reference to specific examples:

[0053] A track defect detection method based on task alignment and dynamic category weighting, such as Figure 1-2 As shown, it includes the following steps:

[0054] S1: Use drones as inspection vehicles, equipped with high-definition camera equipment, to inspect railways or urban tracks and acquire track image data;

[0055] The drone is a quadcopter with a camera resolution of no less than 3840×2160 pixels to ensure the visibility of minute defects on the track surface (such as detachment, wheel burns, indentation, and wear). It features high flight stability and strong maneuverability, enabling it to adapt to the inspection needs of different scenarios along the railway line.

[0056] S2: Preprocess the acquired images and divide them into training set, validation set and test set;

[0057] Before the orbital images acquired by the UAV are input into the detection model, they need to undergo preprocessing to ensure the accuracy of detection and the efficiency of computation. Preprocessing includes steps such as image annotation, image normalization, size scaling, brightness enhancement, and data augmentation.

[0058] A unified annotation system was established for track defects, with the following categories and numbers: 0-Spalling, 1-WheelBurn, 2-Squat, 3-Corrugation. Each target was labeled according to the YOLO text format using the format: class x_center y_center width height. Here, class represents the category, x_center is the center x-coordinate, y_center is the center y-coordinate, width is the width, and height is the height; all coordinates are normalized relative values. Annotations should tightly enclose the defect, with adjacent defects labeled separately. Difficult or uncertain samples were listed separately and not included in the training set.

[0059] Construct a data organization structure that conforms to the general detection framework and configure the data description file (YAML), including the data root directory, training / validation / test subset index, number of categories, and list of category names.

[0060] The original images were resized. Since the resolution of images actually captured by drone cameras is usually much higher than the model's input size, the images were uniformly scaled to a standard input size of 640×640 to fit the input requirements of the deep learning detection model. By standardizing the input size, not only was memory usage and computational complexity reduced, but the consistency of model processing across different images was also ensured.

[0061] S3: In the training process of the YOLO object detection network, task alignment scoring (TAL-lite) is introduced in the positive sample allocation stage.

[0062] In each training iteration, for each ground truth bounding box "g", a set of positive samples is selected from all candidate predictions (all anchors / candidate boxes on multi-scale features). This ensures that the candidates are both class-like in classification and close to the ground truth in regression, thus achieving "task alignment". Unselected candidates participate in classification learning as negative samples.

[0063] S3 includes the following sub-steps:

[0064] S3.1: For a training image, denote all predicted candidate boxes as a set. All truth boxes are sets All categories are set as Each truth box corresponds to a category. ;

[0065] For each candidate box-truth box pair Calculate the location item and the classification item separately; the location item is the geometric overlap. The classification item is aligned to the truth box category, and the calculation formula is as follows:

[0066] ;

[0067] in, Candidate boxes In category The logit value, which is the output value of the YOLO object detection network; Candidate boxes Category confidence level for The function maps logit values ​​to confidence levels [0,1].

[0068] S3.2: To ensure consistency in the "location-classification" dual tasks, for each Define a comprehensive score, and record the score as follows: The expression is:

[0069] ;

[0070] in, The weighting coefficients control the proportion of "geometric alignment" and "category confidence".

[0071] S3.3: For each truth box ,right Sort the candidate boxes in descending order of score and select the top-k candidate boxes with the highest scores as positive candidate samples. If the same candidate box matches multiple ground truth boxes, only the matching relationship with the ground truth box with the highest score is retained.

[0072] This step does not change the network structure or training hyperparameters, but only replaces the sample allocation strategy.

[0073] S4: In track defect detection tasks, class distribution often exhibits severe imbalance. For example, chipping defects account for a high proportion of samples, while burn or wave-wear defects are scarce. Furthermore, the size of different defect targets varies significantly, with some defects occupying only a tiny area in the image, making them difficult to learn effectively using traditional classification loss functions. Under conventional binary cross-entropy (BCE) or Focal Loss frameworks, these samples often contribute very little, leading to insufficient model learning of scarce classes and small targets.

[0074] To address this, a unified scarcity-size joint weighted classification loss, USW, was designed. This introduces dynamically updated class weights and simultaneously characterizes the impact of class scarcity and target size on learning difficulty, thereby improving overall detection accuracy.

[0075] S4 includes the following sub-steps:

[0076] S4.1: For each category The occurrence frequency and average normalized area were calculated separately. To suppress batch jitter, EMA was used, and the occurrence frequency was determined by the category within the most recent iterations. The dynamic smoothing statistic for the selected positive sample is expressed as:

[0077] ;

[0078] in, and The frequencies of occurrence before and after the update are respectively. This represents the current batch of positive samples. These are the weighted smoothing coefficients. ;

[0079] The average normalized area is the category in recent rounds. The dynamic smoothed average of the normalized area of ​​positive samples is expressed as:

[0080] ;

[0081] in, This represents the mean area of ​​the current batch of positive samples.

[0082] This step ensures the smoothness and stability of the category statistics, avoiding training jitter caused by data fluctuations.

[0083] S4.2: Introduce class weights in multi-class classification loss. To simultaneously characterize the impact of class scarcity and target size on learning difficulty, the scarcity and size terms are linearly weighted and fused into class weights. Obtain the category weight function:

[0084] ;

[0085] in, For weighting coefficients, adjust and Impact on the final weight, ; To prevent the occurrence of constants that are divided by zero, , The parameters are used to adjust the influence of the category scarcity term and the target size term on the weights, respectively. ;

[0086] This weighting formula integrates class frequency and target size through linear weighting, avoiding the problem of excessively large or small values ​​that may result from product weighting. Linear weighting avoids extreme value amplification or reduction, resulting in more stable values, smoother training, and ensuring the stability of calculation and the simplicity of engineering implementation.

[0087] S4.3: To maintain consistency between the classification loss magnitude and the baseline, for After normalization and clipping, the expression is:

[0088] ;

[0089] ;

[0090] in, The normalized class weights, and These are the minimum and maximum values ​​of the category weights, respectively. For the final category weight;

[0091] Normalization maintains the average of all weights at 1, thus ensuring that the overall loss scale remains unchanged; then... , Pruning is used to avoid extreme weights that are too large or too small, which can lead to unstable training.

[0092] S4.4: Will A weighted classification loss is constructed by applying the classification supervision term to each positive sample. :

[0093] ;

[0094] in, Positive sample The true category label, Positive samples output by the model The predicted probability of the category. The binary cross-entropy loss function is... Positive sample The true class weights, For the positive sample set, The number of positive samples.

[0095] By employing a joint suppression mechanism for category scarcity and target size, the model can adaptively enhance gradient attention to scarce categories and small target samples during parameter updates, thereby effectively improving the overall performance of track defect detection and steadily increasing the mAP@0.5 and mAP@0.5:0.95 indices.

[0096] S5: Utilize the trained model to perform inference on the newly acquired orbital images from the UAV, outputting detection results including defect category, location, and confidence level, such as... Figure 3 As shown, this ensures that the model can work stably in both experimental environments and actual inspection scenarios.

[0097] During the experimental verification phase, the optimal weight file best.pt obtained from training was loaded into the improved YOLOv12 model, and images from the track defect verification set and test set were input. Figure 4 Figure (a) shows the statistical distribution of track defect samples. The number of samples for the four defect types are: Spalling (983 images), Wheel Burn (382 images), Squat (1265 images), and Corrugation (552 images). Rail Burn samples are the most numerous, accounting for approximately 36%, while Wheel Burn samples are the fewest, at only about 11%, indicating a class imbalance in the dataset. This provides a basis for introducing a class weighting strategy (USW module) in the subsequent model. Figure (b) shows the overlap and distribution of all target boxes. It can be observed that the geometric center of most defect target boxes is concentrated in the middle of the image, and the target size is small, exhibiting a clear clustering phenomenon. This indicates that track defects are usually located in the central region and are small in scale in inspection images, increasing the difficulty of detection. Figures (c) and (d) show the two-dimensional distribution density of the target box center point coordinates (x, y) and width and height, respectively. It can be seen that the target center point is mainly concentrated in the region of x ≈ 0.4 - 0.6 and y ≈ 0.3 - 0.6; while the width and height distribution is mainly concentrated in the region of width < 0.3 and height < 0.3, indicating that most of the orbital defects are small-scale targets.

[0098] The model first performs preprocessing operations such as normalization and resizing on the input image, and then outputs the defect category prediction probability, bounding box coordinates, and confidence score. For example... Figure 5As shown in Figures 6 and 7, the Precision-Confidence, Recall-Confidence, and F1-Confidence curves illustrate the performance changes of the model on the validation set during the training phase. During training, these curves effectively reflect the changing trends of the model's precision, recall, and F1 score at different confidence thresholds, thus demonstrating the model's ability to identify different categories during the learning process. Figure 5 The precision curves are shown at different confidence thresholds. As the confidence level increases, the model's precision across all categories shows a continuous upward trend, especially at a confidence threshold of 0.974, where the average precision for all categories reaches 1.00. This result indicates that the model can effectively identify and classify most defect categories under high confidence conditions. However, some categories, such as corrugation, show lower precision at lower confidence levels, indicating that the model's ability to identify these categories still needs further optimization. Figure 6 shows the trend of recall at different confidence thresholds. It can be observed that in the low confidence range (confidence < 0.2), the model's recall is high, close to 0.96, indicating that the model can identify most defects, especially small, difficult-to-detect defects. However, as the confidence level increases, the recall gradually decreases, indicating that while precision improves, some targets are missed more often, which is closely related to the trade-off between precision and recall. Figure 7 This demonstrates the variation of the F1 score (F1-Score) at different confidence thresholds. The F1 score combines precision and recall, reflecting the model's overall detection capability. According to... Figure 7 As shown, the model achieved its optimal F1 score of 0.69 at a confidence level of 0.335, indicating a good balance between precision and recall, suggesting that the model can achieve the best performance balance at this threshold.

[0099] Finally, a comparative experiment was conducted with the original YOLOv12 model, and the performance evaluation results are as follows:

[0100] Precision: 63.8%;

[0101] Recall rate: 80.4%;

[0102] Average accuracy (mAP@0.5, mAP@0.5:0.95): 73.1% 44.7%;

[0103] F1 score: 71.2%.

[0104] Experimental results show that this invention achieves superior detection accuracy and robustness compared to traditional YOLOv12 series models in detecting small-target defects such as spalling, wheel burns, railhead crushing, and corrugated wear, especially demonstrating a significant advantage in reducing false detection rates in off-track areas. Compared to the original YOLOv12, it improves accuracy by approximately 2.93 percentage points, recall by 9.68 percentage points, and F1 score by 4.23 percentage points, validating the effectiveness of the task alignment and class reweighting strategy. Compared to traditional YOLOv12, this invention primarily addresses the issue of missed detections, resulting in superior overall detection performance.

[0105] Of course, the above description is not intended to limit the present invention, and the present invention is not limited to the examples given above. Any changes, modifications, additions or substitutions made by those skilled in the art within the scope of the present invention should also fall within the protection scope of the present invention.

Claims

1. A method for detecting track defects based on task alignment and dynamic category weighting, characterized in that, Includes the following steps: S1: Use drones equipped with high-definition cameras to inspect railways or urban tracks and acquire track image data; S2: Preprocess the acquired images and divide them into training set, validation set and test set; S3: In the training process of the YOLO object detection network, task alignment scoring is introduced in the positive sample allocation stage; S4: Design a unified scarcity-size joint weighted classification loss, introduce dynamically updated class weights, and simultaneously characterize the impact of class scarcity and target size on learning difficulty, thereby improving overall detection accuracy; S5: Use the trained model to infer the new orbital images acquired by the UAV and output the detection results including the defect category, location and confidence level; S4 includes the following sub-steps: S4.1: For each category Calculate the frequency of occurrence and the average normalized area, respectively. The frequency of occurrence is the category within the most recent iterations. The dynamic smoothing statistic for the selected positive sample is expressed as: ； in, and The frequencies of occurrence before and after the update are respectively. This represents the current batch of positive samples. These are the weighted smoothing coefficients. ; The average normalized area is the category in recent rounds. The dynamic smoothed average of the normalized area of ​​positive samples is expressed as: ； in, This represents the mean area of ​​the current batch of positive samples. S4.2: Combine the scarcity and size terms using a linear weighting to form the class weights. Obtain the category weight function: ； in, For weighting coefficients, adjust and Impact on the final weight, , To prevent the occurrence of constants that divide by zero, , The parameters are used to adjust the influence of the category scarcity term and the target size term on the weights, respectively. ; S4.3: To After normalization and clipping, the expression is: ； ； in, The normalized class weights, and These are the minimum and maximum values ​​of the category weights, respectively. For the final category weight; S4.4: Will A weighted classification loss is constructed by applying the classification supervision term to each positive sample. : ； in, Positive sample The true category label, Positive samples output by the model The predicted probability of the category. The binary cross-entropy loss function is... Positive sample The true class weights, For the positive sample set, The number of positive samples.

2. The orbital defect detection method based on task alignment and dynamic category weighting according to claim 1, characterized in that, In S2, preprocessing includes image annotation, image normalization, size scaling, brightness enhancement, and data augmentation; Rail defects include spalling, wheel burns, rail head crushing, and corrugated wear; Each track defect target is labeled according to the YOLO text format with class, x_center, y_center, width, and height. class is the category, x_center is the center x coordinate, y_center is the center y coordinate, width is the width, and height is the height. All coordinates are normalized relative values.

3. The orbital defect detection method based on task alignment and dynamic category weighting according to claim 1, characterized in that, S3 includes the following sub-steps: S3.1: For a training image, denote all predicted candidate boxes as a set. All truth boxes are sets All categories are set as Each truth box corresponds to a category. ; For each candidate box-truth box pair Calculate the location item and the classification item separately; the location item is the geometric overlap. The classification item is aligned to the truth box category, and the calculation formula is as follows: ； in, Candidate boxes In category The logit value, which is the output value of the YOLO object detection network; Candidate boxes Category Confidence level, for function; S3.2: For each Define a comprehensive score, and record the score as follows: The expression is: ； in, These are weighting coefficients; S3.3: For each truth box ,right Sort the candidate boxes in descending order of score and select the top-k candidate boxes with the highest scores as positive candidate samples. If the same candidate box matches multiple ground truth boxes, only the matching relationship with the ground truth box with the highest score is retained.

Images