Under-vehicle bolt corrosion detection method and device based on improved RT-DETR and medium

Abstract

The invention belongs to the technical field of computer vision and artificial intelligence, and discloses an under-vehicle bolt corrosion detection method and device based on improved RT-DETR and a medium. The method comprises the steps that firstly, an input under-vehicle image is processed by adopting a small target corrosion area enhancement module; the module generates an enhanced feature map capable of highlighting tiny corrosion features through multi-scale feature extraction, a customized corrosion perception attention mechanism and position coding; then, the enhanced feature map is input into an RT-DETR model subjected to lightweight improvement, GhostNet is adopted as a backbone network in the model, and a lightweight neck structure and a corrosion specialization detection head are integrated; and finally, outputting the corrosion grade and the bounding box of the bolt through the model. The method further comprises a targeted small sample data enhancement strategy and a combined loss function based on focus loss, weighted L1 loss and GIoU loss. According to the method, the detection precision of small target corrosion can be remarkably improved, and meanwhile, the model complexity is greatly reduced.

Description

Technical Field

[0001] This invention relates to the fields of computer vision and artificial intelligence, and in particular to a method, equipment and medium for detecting corrosion of undercarriage bolts based on an improved RT-DETR. Background Technology

[0002] With the continuous improvement of the intelligence and informatization of railway transportation equipment, automated and intelligent condition monitoring of key vehicle components has become an important means to ensure driving safety and improve operation and maintenance efficiency. Among them, the undercarriage bolts, as the core components for load-bearing and connection, are subjected to extremely harsh service environments, being exposed to rainwater, salt spray, dust, and electrochemical effects for a long time, making them highly susceptible to corrosion. If corrosion is not detected and treated in time, it may lead to loosening of structural connections, fatigue damage, or even serious safety accidents.

[0003] Currently, the inspection of bolts under trains at railway engineering sites still relies primarily on manual visual inspection. Maintenance personnel must manually observe and subjectively judge thousands of bolts on the train's underside using lighting equipment in maintenance tunnels or depots. This method has significant drawbacks: firstly, it is inefficient, time-consuming, labor-intensive, and extremely demanding; secondly, the inspection results heavily depend on the maintenance personnel's personal experience, sense of responsibility, and mental state, and are prone to omissions and misjudgments under conditions of insufficient lighting, obstructed vision, and complex working environments, resulting in inconsistent inspection quality.

[0004] To overcome the drawbacks of manual inspection, some research institutions and enterprises have begun to explore the introduction of machine vision technology. These attempts mainly fall into two categories: one is to use industrial cameras to acquire images of the undercarriage and then analyze them using traditional image processing algorithms. This type of method typically relies on color space transformation (such as HSV, Lab), edge detection (such as Canny, Sobel), or fixed threshold segmentation to identify rusted areas. However, the environment under a vehicle is extremely complex, with highly uneven lighting distribution, often accompanied by background interference from oil stains, water stains, shadows, and complex pipelines and components. Traditional algorithms are extremely sensitive to these environmental changes. For early rust spots that are small in size and have low contrast with the background, their features are often obscured by noise or background, resulting in poor algorithm robustness and high false positive and false negative rates.

[0005] Another approach employs deep learning, particularly general-purpose object detection networks (such as the YOLO series, Faster R-CNN, and the standard DETR), to identify bolts and their defects. While these methods have achieved great success in general object detection tasks and laid the foundation for the automatic detection of bolt corrosion, they still face significant challenges when applied to the specific scenario of undercarriage bolt corrosion.

[0006] In summary, existing technologies face the following pressing technical challenges when applied to the detection of rust on undercarriage bolts: 1. Challenges in small target recognition: Early and light corrosion usually manifests as tiny rust spots or edge rust marks, occupying only a few pixels in the image. General target detection networks are not capable of extracting features of such small targets, which can easily lead to missed detections or inaccurate detection box localization.

[0007] 2. Training sample scarcity and imbalance: In actual data collection, the number of severely corroded samples is far less than that of normal or lightly corroded samples, resulting in extreme class imbalance. Furthermore, obtaining high-quality labeled samples covering all corrosion morphologies and working conditions is costly, highlighting the sample scarcity problem and limiting the model's generalization ability.

[0008] 3. Limited Model Deployment: Many advanced detection models employ heavy backbone networks such as ResNet and Swin Transformer, resulting in a massive number of model parameters and high computational overhead. This makes them difficult to deploy in real-time on edge computing devices with limited computing power and power consumption, such as vehicle terminals, inspection robots, and trackside monitoring devices, thus failing to balance detection accuracy, inference speed, and engineering application costs.

[0009] 4. Poor adaptability to complex environments: Existing methods are still not adaptable enough to actual working conditions such as complex backgrounds under the vehicle, uneven lighting, and oil stains, and their robustness needs to be improved.

[0010] Therefore, there is an urgent need for a new technology that can effectively solve the four major pain points mentioned above: "difficulty in identifying small targets, scarcity of samples, limited deployment, and class imbalance," and can achieve high-precision, high-real-time detection in complex vehicle undercarriage environments while supporting lightweight deployment. Summary of the Invention

[0011] The main objective of this invention is to overcome the aforementioned deficiencies of the prior art and provide a method, equipment, and medium for detecting corrosion of undercarriage bolts based on an improved RT-DETR, so as to solve the problems of difficulty in identifying small targets, scarcity of training samples, limited model deployment, and imbalance of sample categories in the detection of corrosion of undercarriage bolts.

[0012] In a first aspect, the present invention provides a method for detecting corrosion of under-vehicle bolts based on an improved RT-DETR, comprising the following steps: Acquire images of the area under the vehicle to be inspected; The undercarriage image is processed using a small target rust region enhancement module to generate a final enhanced feature map. The processing steps of the small target rust region enhancement module include: extracting multi-scale features from the undercarriage image to generate a multi-scale feature map; introducing a rust region guiding factor into the multi-scale feature map and using a rust-aware attention mechanism for feature enhancement to obtain an attention-weighted feature map; and fusing the position encoding with the attention-weighted feature map to reduce the dimensionality and obtain the final enhanced feature map. The enhanced feature map is input into the improved RT-DETR model for processing; the improved RT-DETR model uses GhostNet as the backbone network and includes a lightweight neck structure and a corrosion-specialized detection head; The corrosion level and bounding box information of the bolts in the under-vehicle image are output through the corrosion-specialized detection head.

[0013] As an optional implementation of the first aspect of this application, the step of extracting multi-scale features from the under-vehicle image to generate a multi-scale feature map specifically includes: scaling the under-vehicle image to an original resolution scale, a medium resolution scale, and a small resolution scale to obtain three input images of different scales; inputting the three input images of different scales into a shared shallow convolutional network to extract feature maps of three scales; the shallow convolutional network consists of a two-dimensional convolutional layer, a batch normalization layer, and a ReLU activation function; upsampling the feature maps extracted at the medium resolution scale and the small resolution scale so that their size is consistent with the size of the feature map extracted at the original resolution scale; and concatenating the upsampled two feature maps with the feature map extracted at the original resolution scale along the channel dimension to generate the multi-scale feature map.

[0014] As an optional implementation of the first aspect of this application, the steps of introducing a rust region guiding factor into the multi-scale feature map and using a rust-aware attention mechanism for feature enhancement to obtain an attention-weighted feature map specifically include: linearly projecting the multi-scale feature map to generate a query matrix Q, a key matrix K, and a value matrix V; calculating a standard self-attention score based on the query matrix Q and the key matrix K to obtain standard attention weights; introducing a preset rust region guiding factor to weight and enhance the standard attention weights to obtain final attention weights; the rust region guiding factor is an attention map learned and generated during model training to guide the model to focus on rust texture regions; and using the final attention weights to perform a weighted summation of the value matrix V to aggregate and generate the attention-weighted feature map.

[0015] As an optional implementation of the first aspect of this application, the step of fusing the positional encoding with the attention-weighted feature map to obtain the final enhanced feature map specifically includes: constructing a two-dimensional positional encoding matrix based on sine and cosine functions to represent the spatial positional information of pixels in the image; adding the positional encoding to the attention-weighted feature map element by element to obtain a feature map fused with positional information; and using a 1x1 convolutional layer to perform channel compression on the feature map fused with positional information so that its dimension matches the input dimension of the GhostNet backbone network, and finally generating the final enhanced feature map.

[0016] As an optional implementation of the first aspect of this application, before inputting the enhanced feature map into the improved RT-DETR model for processing, a step of preprocessing the undercarriage image using a small-sample rust image data enhancement module is included. The preprocessing step includes at least one of the following: color perturbation enhancement: converting the undercarriage image from the BGR color space to the HSV color space, and applying random perturbations of different ranges to the hue (H) channel, saturation (S) channel, and luminance (V) channel according to four preset levels of normal, light rust, moderate rust, and heavy rust, to generate a color perturbation enhanced feature map; texture simulation overlay: selecting rust spots, cracks, or fissure texture images from a preset rust texture template library, and merging the texture images into the bolt target area in the undercarriage image through transparent overlay to generate a texture overlay enhanced feature map; structural occlusion simulation: randomly sampling oil stains, water stains, or shadow images from a preset occlusion material library, and merging them into the bolt target area in the undercarriage image with preset weights to simulate local occlusion of the undercarriage environment to generate a structural occlusion enhanced feature map.

[0017] As an optional implementation of the first aspect of this application, the processing steps of the lightweight neck structure in the improved RT-DETR model include: channel projection of the feature map output by the GhostNet backbone network using a 1x1 convolutional layer, so that its dimension matches the input dimension of the subsequent Transformer encoder-decoder module, to obtain the projected feature map; and flattening the projected feature map to generate a serialized token sequence, which is used as the input of the Transformer encoder-decoder module.

[0018] As an optional implementation of the first aspect of this application, the processing steps of the corrosion-specialized detection head in the improved RT-DETR model include: reconstructing the target vector output by the Transformer encoder-decoder module to obtain a target vector; sequentially processing the target vector using GhostConv convolution and depthwise separable convolution DWConv to obtain a convolutional feature vector; introducing the corrosion-aware attention mechanism to weight the convolutional feature vector; flattening the weighted feature vector into a one-dimensional vector; and performing classification prediction and bounding box regression through parallel fully connected layers to output the corrosion level and bounding box information.

[0019] As an optional implementation of the first aspect of this application, the method obtains the improved RT-DETR model through model training. The model training steps include: in each round of training, using the Hungarian matching algorithm to establish an optimal one-to-one matching relationship between the model prediction results and the real labels, the matching being based on a matching cost matrix that integrates classification differences, positional biases, and bounding box shape alignment; for successfully matched target pairs, calculating a total loss function to update the model parameters, the total loss function being a weighted sum of classification loss, L1 loss, and GIoU loss; wherein, the classification loss uses a focus loss function to reduce the weight of easily classified background samples and increase the attention to difficult-to-classify rust samples; the L1 loss and the GIoU loss are used together to optimize the center point coordinates, width and height dimensions, and overlap with the real bounding box of the predicted bounding box.

[0020] In a second aspect, embodiments of this application provide an electronic device, which includes a processor, a memory, and a program or instructions stored in the memory and executable on the processor. When the program or instructions are executed by the processor, they implement the steps of the method described in the first aspect.

[0021] Thirdly, embodiments of this application provide a readable storage medium on which a program or instructions are stored, which, when executed by a processor, implement the steps of the method described in the first aspect.

[0022] Compared with the prior art, the beneficial effects of the present invention are as follows: (1) Strong adaptability to small samples: Through the joint data augmentation strategy of color perturbation, texture overlay and occlusion simulation, a large number of virtual samples with diverse shapes and close to real working conditions can be generated online, which effectively solves the problem of scarce and monotonous training samples for bolt corrosion, and significantly improves the robustness and generalization ability of the model in small sample scenarios.

[0023] (2) High accuracy in small target detection: The small target rust area enhancement module is designed to significantly enhance the feature expression of small rust areas by integrating multi-scale feature extraction, customized rust perception attention mechanism and position coding, and suppress the interference of complex background, thereby greatly improving the recognition accuracy of difficult targets such as early rust and small rust spots, and effectively reducing the false detection rate.

[0024] (3) Lightweight and Deployment-Friendly: GhostNet is used as the backbone network, and a lightweight neck structure and corrosion-specific detection head are designed to match it, which greatly reduces the number of model parameters and computational complexity. While ensuring high detection accuracy, the model can achieve high real-time (e.g., 32ms / frame) inference deployment on edge devices with limited power consumption (e.g., ≤15W), which is perfectly adapted to industrial application scenarios such as inspection robots and trackside monitoring terminals.

[0025] (4) Class Imbalanced Adaptation: During the training phase, a combination of Focal Loss, Weighted L1 Loss and GIoU Loss loss functions are used, along with the Hungarian matching algorithm for optimal assignment. This effectively alleviates the class imbalance problem of "more normal bolts and fewer rusted bolts" in the dataset, enabling the model to pay more attention to the difficult-to-classify rusted samples, and improving the ability to distinguish various rusted targets and the accuracy of small target positioning. Attached Figure Description

[0026] Figure 1 This is a flowchart illustrating the overall process framework of the method for detecting corrosion of undercarriage bolts proposed in this invention. Figure 2 This is a comparison chart of the effects of each sub-module in the small sample rust image data enhancement module; Figure 3 This is a schematic diagram of the multi-scale feature extraction network structure of the small target rust region enhancement module in this invention; Figure 4 This is a schematic diagram of the GhostNet backbone network structure used in this invention; Figure 5 This is a schematic diagram of the network structure of the corrosion perception attention mechanism in this invention; Figure 6 This is a schematic diagram of the position encoding and feature projection structure in this invention; Figure 7 This is a schematic diagram of the overall structure of the small target corrosion area enhancement module in this invention; Figure 8 This is a schematic diagram of the rust-specialized detection head in this invention; Figure 9 This is a schematic diagram of the overall framework of the improved RT-DETR model of this invention; Figure 10This is an example of the detection effect diagram and the comparison model effect diagram of the method of the present invention. Detailed Implementation

[0027] The technical solutions of the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of this application, not all embodiments. Based on the embodiments of this application, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this application.

[0028] The terms "first," "second," etc., used in the specification and claims of this application are used to distinguish similar objects and not to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that embodiments of this application can be implemented in orders other than those illustrated or described herein. Furthermore, in the specification and claims, "and / or" indicates at least one of the connected objects, and the character " / " generally indicates that the preceding and following objects are in an "or" relationship. In the description of this invention, "a plurality of" means two or more, unless otherwise explicitly specified.

[0029] Example 1 The core objective of this invention is to address the four major pain points in the detection of rust on under-vehicle bolts: difficulty in identifying small targets, scarcity of samples, limited deployment, and class imbalance. It provides an end-to-end detection solution that balances high precision, real-time performance, and low power consumption.

[0030] This invention proposes a method for detecting corrosion of undercarriage bolts based on an improved RT-DETR, the overall framework of which is as follows: Figure 1 As shown, the framework mainly consists of four interconnected core parts: a small sample corrosion image data augmentation module, a small target corrosion region augmentation module, an improved RT-DETR backbone network, and a training loss function design module.

[0031] The four modules are described in detail below: I. Small Sample Corrosion Image Data Enhancement Module To address the problems of scarce training samples, inconsistent image quality, and diverse corrosion morphologies, this invention designs this module to specifically enhance the original training images. Its structure is as follows: Figure 2 As shown. This module is invoked during the data loading phase and can be conditionally executed based on the corrosion level of each target bolt, enhancing the process to maintain synchronization and consistency between the image and the label. Through this module, data can be... Figure 3 The original image shown in (a) is used to generate the image as shown in (a). Figure 3 The color perturbation image shown in (b) is as follows: Figure 3 The rust texture overlay diagram shown in (c) and as shown in [other diagrams] Figure 3 The structural occlusion map shown in (d) greatly enriches the sample diversity.

[0032] Specifically, this module includes three sub-modules: 1. Enhanced Color Perturbation (Rust Color Shift) First, the original image is converted from BGR space to HSV space. Then, according to the four levels of corrosion characteristics of "normal, light, moderate and heavy", the hue (H), saturation (S) and luminance (V) channels are perturbed in a graded manner. Perturbation formula: Here, clip(˙) represents the perturbation operation. This indicates the hue after the disturbance. This indicates the saturation level after the perturbation. This indicates the brightness after the disturbance. This indicates the applied hue perturbation. This indicates the applied saturation perturbation. This indicates the applied brightness perturbation.

[0033] Finally, the perturbed HSV image is converted back to RGB space to generate a color perturbation enhanced feature map.

[0034] The disturbance parameter settings for different corrosion levels are shown in Table 1.

[0035] Table 1. HSV disturbance range for different corrosion grades 2. Texture Injection To simulate the texture features of rust spots, cracks, and fissures formed by corrosion on the surface of bolts in reality, a real-world rust texture template (containing rust spots, cracks, and fissures) is introduced and fused into the bolt area through transparent overlay.

[0036] Let the bolt area image be The rust texture image is Overlay a transparency mask as The formula for calculating the texture overlay enhancement feature map is: Diverse enhancements can be achieved by adjusting texture size, transparency, and local location.

[0037] 3. Structural Occlusion Simulation Considering the complex interference factors (such as oil stains, water stains, shadows, etc.) in the actual train undercarriage environment, content is randomly sampled from the occlusion material library to simulate the local occlusion of the actual train undercarriage environment, thereby enhancing the model's anti-interference ability and robustness.

[0038] The location, size, and number of occluded areas are randomly generated, and the occluded area does not exceed 30% of the target box area; Occlusion image Within the target region, compare it with the original region by a weight of 0.5. The formula for calculating the fused, structure-occluded enhanced feature map is: The three enhancement submodules mentioned above can be conditionally executed based on the corrosion level of each target bolt, forming a unified multi-source superimposed enhancement function: in For the original image, For a list of tags, To enhance feature maps, this module can be automatically invoked directly in the data loader without modifying the label format, and the enhancement process maintains the synchronization and consistency between the image and the label.

[0039] II. Enhancement Module for Corrosion Areas of Small Targets This module is one of the core components of this invention, designed to enhance the feature representation of small-sized rusted areas at the model input and suppress background interference. Its overall structure is as follows: Figure 4 As shown, it preprocesses the input image before the backbone network.

[0040] 1. Multi-Scale Feature Extraction like Figure 5 As shown, the input image is scaled to three resolution scales, namely the original resolution scale ( ), medium resolution scale ( and small resolution scale ( The formula is as follows: in Represents the original input image (size: ), This indicates bilinear interpolation downsampling of the image. Then, feature extraction is performed. A shallow convolutional network shared by each image input extracts multi-scale corrosion features, as shown in the following formula: in This represents a shared shallow convolutional neural network, consisting of a 2D convolutional layer (Conv2D), a batch normalization layer (BatchNorm2D), and an activation function (ReLU). Each scale corresponds to a feature map. Finally, upsampling is used to restore the original size, and the maps are concatenated along the channel dimension to obtain a multi-scale feature map, as shown in the following formula: in This indicates that bilinear interpolation is used to upsample the smaller feature map back to its original size. This indicates that a splicing operation is performed at the channel level.

[0041] 2. Corrosion perception attention mechanism After obtaining the fused feature map, this invention introduces a three-branch attention structure to generate spatial attention, channel attention, and a local rust guidance mask, enabling the network to automatically focus on local features of the rusted area and improve the recognition ability of small rust spots. For example... Figure 6 As shown, the first step is to use Convolution pairs Linear projection is performed to obtain the query matrix, key matrix, and value matrix transformation matrix. Then, the standard self-attention score is calculated, as shown in the following formula: in Indicates from The The feature vector at each pixel location Indicates passage The query and key transformation matrices generated by convolution. Indicates the first The pixel is the first Attention relevance of pixels, This represents the standard attention weights. Then, to guide the model to focus on the rust texture region, a rust region guiding factor is introduced to enhance attention, resulting in the final attention weighting, as shown in the following formula: in Indicates position The rust-related attention-guiding factors are derived from attention maps learned during training. This represents the final attention weights that incorporate the guidance mechanism. Finally, the attention-weighted feature map is generated by aggregation, as shown in the following formula: in Indicates passage The Value matrix generated by convolution is used to evaluate the input features. Perform a linear transformation.

[0042] 3. Location Encoding and Feature Projection like Figure 7 As shown, since rust often adheres to typical structural locations (such as bolt edges), a sine and cosine position encoding matrix is ​​constructed to make the model spatially sensitive. This matrix is ​​added to the attention features to help the model better understand the positional patterns of bolts and other structures. The position encoding calculation formula is as follows: in Represents the position coordinates of a pixel. Indicates the channel index. This represents the total number of dimensions in the positional encoding. This represents the encoded vector corresponding to the pixel. Adding features can improve the ability to model location. The formula for outputting features after fusing location encodings is: in This is an attention-weighted feature map after corrosion perception weighting, where PE is the sine and cosine positional encoding. This represents the final feature representation that simultaneously incorporates attention weighting and positional encoding.

[0043] To make it consistent with the input dimension of the backbone network, use Convolution compresses the fused features into the dimensions required for the backbone network, as shown in the following formula: in express convolution, This represents the final enhanced feature map, which is then fed into the backbone network.

[0044] III. Improved RT-DETR Backbone Network This part mainly aims to achieve an efficient, low-power, and deployable detection model for deployment of edge devices in industrial scenarios. This invention has made lightweight modifications to the original structure of RT-DETR.

[0045] 1. Lightweight backbone network replacement This invention replaces the original ResNet backbone of RT-DETR with GhostNet, achieving efficient feature representation through the Ghost Bottleneck structure. The backbone network structure is as follows: Figure 8As shown in the diagram, ConvStem is the initial feature extraction module of the GhostNet backbone network. It is responsible for performing preliminary convolution operations on the input features, achieving initial feature compression and dimensionality adjustment, and providing suitable input for subsequent modules. FinalConv is the final convolution module of the GhostNet backbone network, used to integrate the features processed by multiple Ghost Bottleneck modules and output the final features to support subsequent tasks. This network can significantly reduce the number of model parameters and computational cost, making it very suitable for edge computing scenarios.

[0046] 2. Lightweight neck structure Feature map after processing by the enhancement module for the rusted area of ​​the small target. Feature maps are obtained after being fed into the backbone network. To achieve an efficient connection from the backbone feature map to the Transformer structure, this invention introduces a lightweight neck structure after the GhostNet output for channel compression and serialization of the feature map. First, a... The convolutional layer will The required dimensions for projection onto the Transformer are given by the following formula: in The projected feature map is then flattened into a serialized representation, forming the Transformer input token sequence. This is followed by the Transformer encoder-decoder module, which maintains the original RT-DETR structure and takes the projected features from the previous step as input. The output is a candidate set of corrosion targets, as shown in the following formula: Both the encoder and decoder have 6 layers, each layer employing multi-head attention + FFN, and outputting a target vector set. .

[0047] 3. Corrosion-Specialized Detection Head like Figure 9 As shown, this invention improves upon the original detection head of RT-DETR by proposing a corrosion-specialized detection head. First, it uses dimension reconstruction (Reshape) to reshape the target vector. Reconstruct the target vector to obtain the reconstructed target vector. For subsequent operations, then use GhostConv to... Perform phantom convolution to obtain the feature vector after phantom convolution. Then use depthwise separable convolution DWConv to... After processing, the feature vectors after depthwise separable convolution are obtained. To enhance its spatial information modeling capabilities, an attention mechanism for rusted regions is introduced into the small target rusted region enhancement module. The weighted feature vector is obtained by performing weighted processing. The vector is obtained by flattening it. Finally, classification prediction and bounding box regression are performed to obtain the classification probability matrix. and bounding box matrix This outputs the corrosion level and bounding box information, where C is the corrosion type.

[0048] IV. Training Loss Function Design Module During model training, this invention improves and adjusts the original loss function structure of RT-DETR to better adapt to the specific characteristics of rust targets in terms of size, quantity, and category distribution. The Rust Detection Head (Rust-Head) outputs a set of classification results. and bounding box result set The design is mainly achieved through the following function structure: 1. Hungarian matching algorithm During the training phase, the model uses the Hungarian Matcher algorithm to achieve a one-to-one optimal match between the predicted results and the true labels. Specifically, the decoder outputs... A set of candidate targets is composed of several target vectors (each corresponding to a set of classification probabilities and bounding box predictions), while the labeled image contains... The real targets constitute the real set. To establish the optimal matching relationship between them, a... The matching cost matrix, where each element represents the prediction target. With the real target The matching cost is calculated as follows: in, For the first The first prediction target and the second Matching loss between real targets For classifying losses, For the L1 loss of the bounding box, For bounding box shape alignment loss, and The first The predicted category and bounding box of each target. and For the first The label and bounding box of each real target. These represent the weight coefficients for classification loss, bounding box L1 loss, and bounding box shape alignment loss, respectively; the matching loss comprehensively considers three factors: classification difference, positional deviation, and boundary shape alignment. Subsequently, the classic Hungarian algorithm (i.e., the minimum cost solution algorithm for linear assignment problems) is used to find the prediction-true target pair with the minimum total matching loss on this cost matrix, ensuring that the loss function during training achieves optimal allocation at both the semantic and spatial levels, thereby improving training stability and convergence speed.

[0049] 2. Focus Loss Function Based on the optimal matching relationship described above, supervised learning is performed on the matched target pairs. For classification loss, Focal Loss is used instead of the original cross-entropy loss to mitigate the impact of class imbalance. Rust targets are usually in a minority class in the samples and are easily masked by a large amount of background information. Focal Loss effectively improves the model's ability to distinguish minority targets (such as pitting and edge rust) by increasing the weight of hard-to-classify samples. The expression for Focal Loss is as follows: in, To output the predicted classification result value for the detection head, It is a category balance factor. This is the focus parameter (usually set to 2.0).

[0050] 3. Weighted Loss Function For bounding box regression, a weighted combination of L1 loss and GIoU loss is used to jointly optimize the positional accuracy and shape overlap of the predicted box. Specifically, the L1 loss is used to regress the center point coordinates and width / height dimensions. GIoU loss focuses more on the spatial alignment between the predicted bounding box and the ground truth bounding box, making it suitable for handling the problem of small rusted targets being sensitive to IoU changes due to their small scale. in It is the true value of the sample. It is the predicted value of the bounding box output by the detection head.

[0051] Ultimately, the total loss function proposed in this invention consists of the above three terms, and is defined as follows: in, These are the weighting coefficients for each loss term, with default settings of 0.2, 0.5, and 0.2, used to enhance the fitting ability for small target location prediction.

[0052] In this embodiment, the model training and model deployment and detection need to be explained in detail as follows: 1. Model Training Data preparation: High-definition industrial cameras were used to photograph the undercarriage of trains at fixed points in the railway depot, collecting 2457 raw images with a resolution of 640×640. Professionals labeled the bolts in the images, classifying them into four levels: "normal," "light corrosion," "moderate corrosion," and "severe corrosion," generating YOLO format label files.

[0053] Data augmentation: Before training, the original images were augmented by: (1) HSV color perturbation based on the parameters in Table 1; (2) rust texture overlay with transparency α randomly varying between 0.3 and 0.7; and (3) structural occlusion simulation with an occlusion area not exceeding 30% of the target box. These operations expanded the original sample set, generating diverse training data.

[0054] Model Training: The enhanced image is fed into a Small Target Corrosion Region Enhancement Module (SRAM) for processing, yielding the final enhanced feature map. This feature map is then input into an improved RT-DETR model based on GhostNet. The model employs a 6-layer Transformer encoder and a 6-layer decoder. The AdamW optimizer is used, and end-to-end training is performed using the aforementioned weighted combined loss function. The model is trained for 150 epochs, converging and achieving optimal performance on the 95th epoch. This weight file is then saved as the deployment model.

[0055] 2. Model Deployment and Detection Deployment environment: The NVIDIA Jetson NX embedded AI platform was selected as the edge computing device, and its power consumption mode was set to 15W.

[0056] Image Acquisition: The Jetson NX platform with the deployed model is integrated into the track inspection robot. The industrial camera on the robot acquires real-time video streams from under the vehicle and transmits 640×640 resolution image frames to the Jetson NX.

[0057] Real-time detection: After normalization, the received image frames are sequentially processed through the complete detection workflow designed in this invention: SRAM module, GhostNet backbone, lightweight neck, Transformer codec, and rust-specialized detection head. For dense bolt scenes, non-maximum suppression (NMS) is enabled to filter out overlapping detection boxes. Finally, the model outputs the category (rust level) and precise bounding box coordinates of each detected bolt in the image.

[0058] Detection Results: Field tests showed that the method of this invention achieved an average inference speed of 32 milliseconds / frame on the Jetson NX platform, fully meeting the real-time detection requirement of 30 FPS. Figure 10 As shown, compared with the model before the improvement ( Figure 10 (b) and (d)), the method of the present invention ( Figure 10 (a) and (c) can more accurately identify minute rust and partially obscured bolts, significantly reducing missed and false detections.

[0059] Optionally, embodiments of this application also provide an electronic device, including a processor, a memory, and a program or instructions stored in the memory and executable on the processor. When the program or instructions are executed by the processor, they implement the various processes of the above-described embodiment of the method for detecting rust on under-vehicle bolts based on improved RT-DETR, and achieve the same technical effect. To avoid repetition, they will not be described again here.

[0060] This application also provides a readable storage medium storing a program or instructions. When the program or instructions are executed by a processor, they implement the various processes of the above-described embodiment of the method for detecting rust on under-vehicle bolts based on an improved RT-DETR, and achieve the same technical effect. To avoid repetition, these will not be described again here.

[0061] The processor is the processor in the electronic device described in the above embodiments. The readable storage medium includes computer-readable storage media, such as computer read-only memory (ROM), random access memory (RAM), magnetic disk, or optical disk.

[0062] It should be noted that, in this document, the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Without further limitations, an element defined by the phrase "comprising one..." does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes that element. Furthermore, it should be noted that the scope of the methods and apparatuses in the embodiments of this application is not limited to performing functions in the order shown or discussed, but may also include performing functions substantially simultaneously or in the reverse order, depending on the functions involved. For example, the described methods may be performed in a different order than described, and various steps may be added, omitted, or combined. Additionally, features described with reference to certain examples may be combined in other examples.

[0063] Through the above description of the embodiments, those skilled in the art can clearly understand that the methods of the above embodiments can be implemented by means of software plus necessary general-purpose hardware platforms. Of course, they can also be implemented by hardware, but in many cases the former is a better implementation method. Based on this understanding, the technical solution of this application, in essence, or the part that contributes to the prior art, can be embodied in the form of a software product. This computer software product is stored in a storage medium (such as ROM / RAM, magnetic disk, optical disk) and includes several instructions to cause a terminal (which may be a mobile phone, computer, server, air conditioner, or network device, etc.) to execute the methods described in the various embodiments of this application.

[0064] The embodiments of this application have been described above with reference to the accompanying drawings. However, this application is not limited to the specific embodiments described above. The specific embodiments described above are merely illustrative and not restrictive. Those skilled in the art can make many other forms under the guidance of this application without departing from the spirit and scope of the claims, and all of these forms are within the protection scope of this application.

Claims

1. A method for detecting corrosion of undercarriage bolts based on an improved RT-DETR, characterized in that, Includes the following steps: Acquire images of the area under the vehicle to be inspected; The undercarriage image is processed using a small target rust region enhancement module to generate a final enhanced feature map; The processing steps of the small target rust region enhancement module include: extracting multi-scale features from the under-vehicle image to generate a multi-scale feature map; introducing a rust region guiding factor into the multi-scale feature map and using a rust-aware attention mechanism to enhance the features, resulting in an attention-weighted feature map; and fusing the position encoding with the attention-weighted feature map to reduce the dimensionality and obtain the final enhanced feature map. The enhanced feature map is input into the improved RT-DETR model for processing; the improved RT-DETR model uses GhostNet as the backbone network and includes a lightweight neck structure and a corrosion-specialized detection head; The corrosion level and bounding box information of the bolts in the under-vehicle image are output through the corrosion-specialized detection head.

2. The method according to claim 1, characterized in that, The steps of extracting multi-scale features from the under-vehicle image to generate a multi-scale feature map specifically include: The under-vehicle image is scaled to the original resolution scale, the medium resolution scale, and the small resolution scale to obtain three input images at different scales; The three input images at different scales are respectively input into a shared shallow convolutional network to extract feature maps at three scales; the shallow convolutional network consists of a two-dimensional convolutional layer, a batch normalization layer, and a ReLU activation function. The feature maps extracted at the medium and low resolution scales are upsampled to make their size consistent with the feature map size extracted at the original resolution scale. The upsampled two feature maps are concatenated with the feature map extracted at the original resolution scale along the channel dimension to generate the multi-scale feature map.

3. The method according to claim 2, characterized in that, The steps of introducing a corrosion region guiding factor into the multi-scale feature map and using a corrosion-aware attention mechanism for feature enhancement to obtain an attention-weighted feature map specifically include: Linear projection is performed on the multi-scale feature map to generate a query matrix Q, a key matrix K, and a value matrix V; Based on the query matrix Q and the key matrix K, a standard self-attention score is calculated to obtain the standard attention weight; A preset rust region guiding factor is introduced to weight and enhance the standard attention weights to obtain the final attention weights; the rust region guiding factor is an attention map learned and generated during model training to guide the model to focus on rust texture regions; The value matrix V is weighted and summed using the final attention weights to generate the attention-weighted feature map.

4. The method according to claim 3, characterized in that, The step of fusing the positional encoding with the attention-weighted feature map to obtain the final enhanced feature map specifically includes: A two-dimensional position coding matrix based on sine and cosine functions is constructed to represent the spatial position information of pixels in an image; The location code is added element-wise to the attention-weighted feature map to obtain a feature map that incorporates location information; A 1x1 convolutional layer is used to perform channel compression on the feature map that incorporates location information, so that its dimension matches the input dimension of the GhostNet backbone network, and finally the final enhanced feature map is generated.

5. The method according to claim 1, characterized in that, Before inputting the enhanced feature map into the improved RT-DETR model for processing, the method further includes a step of preprocessing the undercarriage image using a small-sample rust image data enhancement module. This preprocessing step includes at least one of the following: Color perturbation enhancement: The under-vehicle image is converted from the BGR color space to the HSV color space. Based on four preset levels of normal, light corrosion, moderate corrosion, and heavy corrosion, different ranges of random perturbation are applied to the hue (H) channel, saturation (S) channel, and luminance (V) channel to generate a color perturbation enhancement feature map. Texture simulation overlay: Select rust spots, cracks or fissures texture images from the preset rust texture template library, and fuse the texture images into the bolt target area in the undercarriage image through transparent overlay to generate a texture overlay enhancement feature map; Structural occlusion simulation: Randomly sample oil stains, water stains, or shadow images from a preset occlusion material library, and fuse them with preset weights into the bolt target area in the under-vehicle image to simulate local occlusion of the under-vehicle environment and generate a structural occlusion enhancement feature map.

6. The method according to claim 1, characterized in that, The processing steps for the lightweight neck structure in the improved RT-DETR model include: A 1x1 convolutional layer is used to perform channel projection on the feature map output by the GhostNet backbone network, so that its dimension matches the input dimension of the subsequent Transformer encoder-decoder module, thus obtaining the projected feature map. The projected feature map is flattened to generate a serialized token sequence, which serves as the input to the Transformer encoder-decoder module.

7. The method according to claim 1 or 6, characterized in that, The processing steps of the corrosion-specialized detection head in the improved RT-DETR model include: The target vector is reconstructed from the output of the Transformer encoder-decoder module to obtain the target vector; The target vector is processed sequentially using GhostConv convolution and depthwise separable convolution DWConv to obtain the convolutional feature vector. The corrosion-sensing attention mechanism is introduced to weight the convolutional feature vector; Flatten the weighted eigenvectors into one-dimensional vectors; The corrosion level and bounding box information are output by performing classification prediction and bounding box regression separately through parallel fully connected layers.

8. The method according to claim 1, characterized in that, The method obtains the improved RT-DETR model through model training, and the model training steps include: In each round of training, the Hungarian matching algorithm is used to establish the optimal one-to-one matching relationship between the model prediction results and the real labels. This matching is based on a matching cost matrix that integrates classification difference, positional bias and bounding box shape alignment. For a successfully matched target pair, a total loss function is calculated to update the model parameters. The total loss function is a weighted sum of the classification loss, L1 loss, and GIoU loss. The classification loss uses a focus loss function to reduce the weight of easily classified background samples and increase the attention to difficult-to-classify rust samples. The L1 loss and the GIoU loss are used together to optimize the center point coordinates, width and height dimensions, and overlap with the true bounding box of the predicted bounding box.

9. An electronic device, characterized in that, It includes a processor, a memory, and a program or instructions stored in the memory and executable on the processor, wherein when the program or instructions are executed by the processor, they implement the steps of a method for detecting corrosion of under-vehicle bolts based on an improved RT-DETR as described in any one of claims 1-8.

10. A readable storage medium, characterized in that, The readable storage medium stores a program or instructions, which, when executed by a processor, implement the steps of the method for detecting corrosion of under-vehicle bolts based on an improved RT-DETR as described in any one of claims 1-8.

Images