A steel rail surface defect detection method based on a diffusion model drive

Abstract

The present application belongs to the field of track defect detection, and discloses a steel rail surface defect detection method based on a diffusion model drive. In view of the problem of lack of industrial defect annotation data, a training framework based on dynamic defect generation and semi-supervised is proposed, a small amount of real annotation data is used to synthesize a large number of high-fidelity defect samples; by simulating the irregularity of the shape of the real defect, the edge transition and the material characteristics, the physical authenticity and the boundary naturalness of the defect generation are improved; through a multi-level fusion strategy, the seamless fusion of the illumination, texture and color of the generated defect and the original background is realized, and the visual consistency of the generated defect and the background is realized; through dynamic pseudo-label generation and difficult sample focusing strategy, the small amount of labeled data and a large amount of unlabeled generated data are maximized, and efficient training under the condition of insufficient labeled data is realized; by introducing an adversarial training mechanism, the generator learns the detail distribution of the real defect, and the boundary accuracy and structural coherence of the segmentation mask are improved.

Description

Technical Field

[0001] This invention belongs to the field of rail defect detection and relates to a method for detecting rail surface defects based on a diffusion model. Background Technology

[0002] With the rapid development of high-speed railways and urban rail transit in my country, the surface condition of rails, as the core load-bearing component of the track system, directly affects train operation safety. During long-term service, rails can develop various defects such as cracks, corrosion, and spalling. If these defects are not detected and addressed in a timely manner, they can lead to serious safety accidents. Traditional inspection methods mainly rely on manual inspection and image processing technology, which suffer from low efficiency and high missed detection rates. Manual inspection is not only time-consuming and labor-intensive but also easily affected by human factors, making it difficult to guarantee consistency and accuracy. While automated inspection systems based on traditional image processing technology have improved efficiency to some extent, their accuracy in complex environments remains insufficient.

[0003] Currently, mainstream automated inspection systems primarily rely on fully supervised deep learning-based models, which heavily depend on large amounts of high-quality labeled data for training. However, obtaining high-precision pixel-level annotations (masks) of rail surface defects is costly, time-consuming, and requires a high level of expertise, resulting in extremely limited labeled samples available for training. Secondly, in practical industrial applications, unlabeled rail image data often far exceeds labeled data; therefore, effectively utilizing this unlabeled data to improve model performance is a major challenge in current research. The limitations of fully supervised learning methods directly affect the model's generalization ability in real-world industrial scenarios, making it difficult for the model to cope with complex and ever-changing real-world environments. Furthermore, rail defects are diverse in type, shape, and boundary (e.g., the ductility of cracks and the irregularity of spalling). Existing models struggle to effectively learn these highly dynamic defect features under small sample conditions and have poor adaptability to newly emerging defect types or environmental changes.

[0004] While generative models such as Generative Adversarial Networks (GANs) have made some progress in image generation, they still suffer from problems such as mode collapse and image distortion. GANs are prone to instability during training, and the generated images may contain unrealistic or physically inconsistent features. Although diffusion models have made breakthroughs in image generation quality, exhibiting better stability and diversity, their optimization for specific industrial inspection scenarios remains insufficient. Therefore, there is an urgent need to propose a robust rail surface defect detection technology that can effectively overcome the scarcity of high-quality labeled data, accurately simulate the complex and varied morphology of rail defects, achieve natural integration of defects with the background, and efficiently utilize massive amounts of unlabeled data. This would improve the accuracy, efficiency, and generalization ability of detection, meeting the urgent needs of safe operation and maintenance in the rapidly developing rail transit industry. Summary of the Invention

[0005] The purpose of this invention is to propose a rail surface defect detection method based on diffusion model. This method integrates dynamic mask enhancement, diffusion model generation and semi-supervised learning techniques to process rail surface defects, so as to solve the technical problems of scarce track defect data, high cost of manual annotation and insufficient defect detection accuracy in the existing technology.

[0006] To achieve the above objectives, the present invention adopts the following technical solution: A method for detecting rail surface defects based on a diffusion model includes the following steps: Step 1. Acquire the original noise-free image of the rail surface and generate the original defect mask through manual annotation. The original noise-free image and the corresponding original defect mask together form the rail defect dataset. For the original defect mask, dynamic mask generation is performed based on multi-level morphological transformation to obtain an enhanced defect mask; Step 2. For the original noise-free image and the enhanced defect mask, an improved conditional diffusion model architecture is adopted. In the latent space of the image, the enhanced defect mask is used as a condition, and the diffusion process is guided by dual-path control to generate a noisy image with a fused preset defect shape. After denoising, the defect image is obtained. Step 3. For the defect area of ​​the generated defect image and the background area of ​​the original noise-free image, a three-level fusion strategy is used to fuse the defect area into the background area to obtain a visually highly consistent fused image; Step 4. Build a segmentation network and pre-train it using an industrial defect dataset to obtain a pre-trained defect segmentation model; Step 5. Using an adversarial training framework, the pre-trained defect segmentation model is used as the generator in the generative adversarial network. The generator is fine-tuned using the rail defect dataset to obtain an optimized student generator and a teacher model that is stable through EMA. Step 6. Based on the stable teacher model of EMA, pseudo-labels are generated for the fused image, and difficult samples are identified through confidence screening, feature clustering, and dynamic reweighting. After data augmentation is applied to the difficult samples, they are input into the student generator, i.e., the student model, for training. The teacher model updates its parameters smoothly through EMA, and the optimal model weights are obtained through training. By using a trained student model with optimal model weights, the surface defects of the rail can be segmented.

[0007] Furthermore, based on the aforementioned diffusion model-driven rail surface defect detection method, this invention also proposes a corresponding diffusion model-driven rail surface defect detection system, which adopts the following technical solution: A rail surface defect detection system based on a diffusion model includes the following modules: The data acquisition and dynamic enhancement processing module is used to acquire the original noise-free image of the rail surface and generate the original defect mask through manual annotation. The original noise-free image and the corresponding original defect mask together form the rail defect dataset. For the original defect mask, dynamic mask generation is performed based on multi-level morphological transformation to obtain an enhanced defect mask; The defect generation module is designed to generate a noisy image with a fused preset defect shape by using an improved conditional diffusion model architecture for the original noise-free image and the enhanced defect mask. In the latent space of the image, the enhanced defect mask is used as a condition and the diffusion process is guided by dual-path control to generate a noisy image with a fused preset defect shape. After denoising, the defect image is obtained. The image seamless fusion processing module is used to fuse the defective region of the generated defective image into the background region of the original noise-free image using a three-level fusion strategy, so as to obtain a visually highly consistent fused image. The defect segmentation model building module is used to build a segmentation network including a hybrid encoder and a decoder; the segmentation network is pre-trained on the MVTec dataset to obtain a pre-trained defect segmentation model; The supervised training module is used to employ an adversarial training framework, using the defect segmentation model as the generator in the generative adversarial network. The generator is fine-tuned using a rail defect dataset to obtain an optimized student generator and a teacher model that is stable through EMA. The semi-supervised training module is used to generate pseudo-labels for the fused images based on the EMA-stable teacher model. The hard samples are then identified through confidence screening, feature clustering, and dynamic reweighting. After data augmentation is applied to the hard samples, they are input into the student generator, i.e., the student model, for training. The teacher model updates its parameters smoothly through EMA, and finally obtains the optimal model weights. By using a trained student model with optimal model weights, the surface defects of the rail can be segmented.

[0008] Furthermore, based on the aforementioned diffusion model-driven rail surface defect detection method, this invention also proposes a computer device comprising a memory and one or more processors. Executable code is stored in the memory. When the processor executes the executable code, it implements the steps of the aforementioned diffusion model-driven rail surface defect detection method.

[0009] Furthermore, based on the above-mentioned diffusion model-driven rail surface defect detection method, this invention also proposes a computer-readable storage medium storing a program that, when executed by a processor, implements the steps of the above-mentioned diffusion model-driven rail surface defect detection method.

[0010] The present invention has the following advantages: As described above, this invention presents a method for detecting rail surface defects based on a diffusion model. Addressing the scarcity of labeled industrial defect data, this method proposes a combination of dynamic defect generation and a semi-supervised training framework. This enables the synthesis of a massive number of high-fidelity defect samples using a small amount of real labeled data, solving the problem of sample scarcity. Secondly, by simulating the irregularity of the morphology, edge transitions, and material characteristics of real defects, this invention improves the physical realism and naturalness of the generated defects' boundaries, avoiding the harsh boundaries of traditional methods. Thirdly, this invention achieves seamless integration of the generated defects with the original background's lighting, texture, and color through a multi-level fusion strategy, achieving visual consistency between the generated defects and the background and avoiding traces of artificial synthesis. Fourthly, this invention maximizes the use of a small amount of labeled data and a large amount of unlabeled generated data through dynamic pseudo-label generation and hard sample focusing strategies, achieving efficient training under conditions of insufficient labeled data and solving the problem of high labeling costs in industrial scenarios. Fifthly, this invention introduces an adversarial training mechanism to drive the generator to learn the detailed distribution of real defects. Through discriminator supervision, it forces the segmentation mask output by the generator to approximate the real defect label in terms of structural coherence and boundary accuracy, improving the segmentation accuracy of defects. The method of this invention integrates dynamic mask enhancement, diffusion model generation and semi-supervised learning to process rail surface defects, effectively solving the technical problems of scarce track defect data, high cost of manual annotation and insufficient defect detection accuracy in the prior art. Attached Figure Description

[0011] Figure 1 This is a flowchart of the rail surface defect detection method based on diffusion model driven in an embodiment of the present invention; Figure 2 This is a schematic diagram of the edge control module in an embodiment of the present invention; Figure 3 This is a schematic diagram of the dual-channel control module in an embodiment of the present invention; Figure 4 This is a schematic diagram of the UNet architecture in an embodiment of the present invention; Figure 5 This is a comparison image between the fused image generated based on the three-level fusion strategy in an embodiment of the present invention and the original image; Figure 6 This is a diagram of the segmentation network structure built based on the Mamba module, Transformer, and CBAM in an embodiment of the present invention; Figure 7 This is a structural block diagram of the generative adversarial network constructed in an embodiment of the present invention; Figure 8 This is an image showing the effect of segmenting defects on the surface of the rail after applying the method of the present invention. Detailed Implementation

[0012] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments: Example 1 like Figure 1 As shown in the figure, this embodiment describes a rail surface defect detection method based on diffusion model driving. It integrates dynamic mask enhancement, diffusion model generation and semi-supervised learning techniques to process rail surface defects, so as to solve the technical problems of scarce track defect data, high cost of manual annotation and insufficient defect detection accuracy in the existing technology.

[0013] Specifically, the rail surface defect detection method based on diffusion model driven in this embodiment includes the following steps: Step 1. Data acquisition and dynamic enhancement processing.

[0014] The original noise-free image of the rail surface is acquired, and the original defect mask is generated by manual annotation. The original noise-free image and the corresponding original defect mask are combined to form a rail defect dataset.

[0015] For the original defect mask, dynamic mask generation is performed based on multi-level morphological transformation to obtain an enhanced defect mask.

[0016] Specifically, this invention addresses the problems of scarce samples and varied defect morphologies in industrial defect detection by using an industrial camera array to acquire raw, noise-free images of the railway rail surface and manually annotating the raw defect mask.

[0017] Then, dynamic mask generation is performed based on multi-level morphological transformation. Random morphological operations, multi-scale Gaussian noise injection, and random rotation transformation are sequentially applied to the original defect mask to simulate the morphological variability and irregular boundaries of real defects.

[0018] Through the above multi-level morphological transformations, a dynamically generated enhanced defect mask with continuous edge transitions and spatial diversity can be obtained. The original noise-free image and the enhanced defect mask are the key inputs for the defect image generation in step 2.

[0019] The method of this invention uses a real defect mask as a base and employs a three-level cascaded enhancement strategy, as shown below: 1. Morphological topological transformation layer: First, the original noise-free image is obtained through manual annotation. The mask is then applied to the original noise-free image. The mask is subjected to a random morphological operation sequence, including dilation, erosion, opening, closing, etc.

[0020] Each operation kernel randomly selects an elliptical or rectangular structure, and the kernel size is randomly generated following a uniform distribution. Mathematically, this is expressed as: .

[0021] in, This is the mask after morphological transformation. Let i = 1 to N be a cascaded sequence of operations, representing a chained combination of multiple morphological operations. For the first Submorphological operations. This is the original defect mask.

[0022] 2. Multi-scale noise fusion layer: Random noise is injected across multiple scales to simulate the irregular boundaries of real defects.

[0023] .

[0024] in, For the mask after injecting noise, The noise scale, i.e., the resolution level, includes 4×4, 8×8, and 16×16. It is Gaussian noise with a mean of 0 and a variance of . . To reduce noise through bilinear interpolation Adjust to the original image size, height H, width W.

[0025] 3. Geometric Deformation Layer: Applying random rotation transformations to enhance spatial diversity, rotation angle The matrix form is: .

[0026] in, for The original coordinates of the middle pixel The coordinates are the transformed coordinates. The rotation angle follows a uniform distribution. - , , This is the translation amount, to ensure rotation around the image center.

[0027] Through multi-level morphological transformations, enhanced defect masks with continuous edge transitions can be generated, effectively solving the problem of harsh boundaries caused by traditional random rectangular masks, and making the generated defects more in line with the physical diffusion law.

[0028] Step 2. Construction of the diffusion model-driven defect generation module.

[0029] For the original noise-free image and the enhanced defect mask, an improved conditional diffusion model (ControlNet-Enhanced Stable Diffusion) architecture is adopted. In the image latent space, the enhanced defect mask is used as a condition, and the diffusion process is guided by dual-path control to generate a noisy image with fused preset defect shapes. This noisy image is then denoised to obtain the defect image. The generated defect image and the corresponding original background region (i.e., The background is the object of the seamless integration process in step 3.

[0030] Step 2 employs an improved conditional diffusion model architecture to achieve high-quality defect generation in the latent space.

[0031] Step 2 mainly includes the following parts: 1. Define the latent variable diffusion process, and set the cumulative noise coefficient through the noise addition formula and cosine scheduler to retain more signal in the early stage of training, so as to improve convergence stability and obtain noisy image.

[0032] Latent variables are obtained through Figure 4 The VAE encoding mentioned above can compress large images into small latent spaces, reducing computation while preserving key features. Noise is added to the latent variables according to random time steps t, and then the model learns to predict the noise added to the low-dimensional latent space. A clear image is gradually generated through subsequent iterations to remove noise.

[0033] Define the latent variable diffusion process, and the noise addition formula is: .

[0034] in, For a noisy image at time step t, The original noise-free image, It is standard Gaussian noise.

[0035] The cumulative noise figure, used to control the degree of noise addition, is set according to the cosine tuner, where: .

[0036] Where T is the total number of time steps. This is the initial time step offset.

[0037] This cosine scheduler can retain more signal in the early stages of training, which helps improve convergence stability.

[0038] This embodiment expands the exploration range of the latent space through a noise offset mechanism.

[0039] 2. For example Figure 4 As shown, the diffusion model adopts a symmetrical UNet architecture, with the encoder on the left, which consists of multiple coding blocks, the bottleneck layer in the middle, and the decoder on the right, which consists of multiple decoding blocks.

[0040] Each coding block consists of a corresponding convolutional layer, a dual-path control module, and a downsampling layer (implemented using pooling operations). Its purpose is to gradually extract high-level features of the image and compress spatial resolution through downsampling.

[0041] The bottleneck layer consists of convolutional layers, which further extract high-level features and minimize the size of the feature maps.

[0042] The dual-path control module is a component of the UNet's encoding and intermediate blocks, used to guide generation. Furthermore, this invention also includes an edge control module and a material control module, such as... Figure 3 As shown. The outputs of the edge control module and the material control module serve as the inputs of the dual-channel control module, which is used to fuse the two types of information.

[0043] The edge control module extracts edge features from the original noise-free image using the Canny operator, outputs a feature map after processing through multiple residual blocks, and then downsamples it to serve as input to each coding block and intermediate block of the diffusion model UNet.

[0044] The edge control module first extracts the original noise-free image using the Canny edge detection operator. The edge features in the image are calculated using the following formula: Canny is an edge detection operator.

[0045] in, For edge feature maps, For the high threshold of the Canny operator, This is the low threshold of the Canny operator.

[0046] like Figure 2 As shown, the extracted edge features After feature extraction and channel adjustment of three residual blocks, the final feature map is output. This output feature map is then downsampled and used sequentially as input to encoding blocks 2, 3, and 4, as well as the intermediate blocks, of the diffusion model UNet. Figure 4 As shown, this guides the diffusion process and prevents the generated content from being structurally unreasonable.

[0047] The material control module is used to receive material text input and obtain text embedding features through a text encoder.

[0048] Material information, as part of the text prompt, is encoded by the text encoder and then cross-attentioned with the intermediate features of the UNet's encoding blocks. This results in a feature map containing material information, which guides the diffusion process. Figure 3 As shown.

[0049] Material information text, such as descriptions related to materials like "rail surface," "metallic texture," "rust deposits," and "wear marks," is generated. This material description is converted into vectors (text embedding features) by a text encoder. These vectors are then cross-attentioned with the intermediate features of the encoded block, outputting a feature map containing material information. The material control module here is designed to better control and guide the image generation process, making the generated defect images more realistic and believable in both visual and material characteristics.

[0050] Dual-channel control modules are set at each coding block and intermediate block of UNet. The input of the dual-channel control modules are the output features of the edge control module after downsampling and the text embedding features output by the material control module.

[0051] The dual-path control module performs cross-attention on the output features of the downsampled edge control module and the text embedding features output by the material control module, and outputs a feature map with material information, thereby guiding the diffusion process.

[0052] like Figure 3 The diagram shows the structure of each network module of the dual-channel control module, and its internal processing flow is as follows: The material control module receives the material text input and obtains the text embedding features through a text encoder, serving as one path of the dual-path control module. The edge control module extracts edge features from the original noise-free image using the Canny operator, then further extracts features through residual blocks. After downsampling to adapt to the sizes of different UNet coding blocks, the output serves as the other path of the dual-path control module. The two outputs are then fused with material text information and edge image information through cross-attention within the dual-path control module. Finally, through operations such as convolution, a feature map containing both material and edge information is output.

[0053] The decoder section progressively restores the spatial resolution of the image through upsampling layers. The goal of the decoder is to map the feature information extracted by the encoder back to the original image size. In each layer of the encoder, skip connections directly pass low-level feature maps to the corresponding decoder layer, aiming to preserve high-resolution spatial information.

[0054] UNet, as the core component, predicts and removes noise at multiple time steps, ultimately generating high-quality defect images.

[0055] 3. By jointly optimizing the noise prediction loss, perception loss, and contrast loss, the diffusion model generates images that are both realistic and effectively highlight defect areas, while maintaining consistency in the noise distribution of the latent space.

[0056] Total loss function of diffusion model The calculation formula is: ; ; ; .

[0057] in, To constrain the consistency of potential spatial noise distribution for noise prediction loss; To mitigate perceptual loss, a pre-trained VGG-19 network is used to maintain high-level feature similarity to ensure that the generated image is perceptually similar to the real image. This is used to reduce contrast and enhance the contrast between the defective area and the background, making the defect more prominent. , , These are the weighting coefficients for each loss term. These coefficients are used to balance the effects of noise prediction, perceptual similarity, and contrast. Through training, these coefficients can be adaptively adjusted to optimize model performance. Noise predicted by UNet; These are the parameters of the network; The image is a noisy image at time step t, where t is a random time step. c represents the actual noise; c represents the conditional input, namely the enhanced defect mask, edge feature map, and text embedding features. Indicates the first The number of channels in the layer feature map; Indicates the first The height of the layer feature map; Indicates the first The width of the layer feature map; Indicates the first Layer features; Defect images generated for the diffusion model; The original, noise-free image; The feature vector of the defect region; Background features; τ is the feature vector of the k-th region, used to calculate the contrast between the background and the defect; τ is the temperature coefficient, which controls the smoothness of the probability distribution. This indicates the total number of regions.

[0058] Through joint optimization of the above loss functions, the diffusion model can generate images that are both realistic (through perceptual loss) and effectively highlight defective areas (through contrast loss), while maintaining consistency in the noise distribution of the latent space.

[0059] This invention solves the problem of harsh boundaries in traditional rectangular masks by generating dynamic masks in step 1 above, and simulates the physical diffusion law (continuous edge transition). By using the dual-path control (edge ​​+ material) of the diffusion model in step 2, it ensures that the defect texture is consistent with the background. The joint loss function enhances the saliency of defects, thus significantly improving the realism of defect generation.

[0060] Step 3. Seamless image fusion processing.

[0061] For the defect image generated in step 2 The defect area compared to the original noise-free image The background area employs a three-level fusion strategy to... The defective areas are blended into the background area, resulting in a visually highly consistent fused image. .

[0062] Step 3.1. Gradient domain Poisson fusion.

[0063] Using the background of the original noise-free image as a reference, the defect region of the defect image generated in step 2 is fused to the background based on gradient domain Poisson fusion, minimizing gradient differences and boundary inconsistencies, as shown in the following formula: .

[0064] in, The original noise-free image, B is the defect image generated in step 2, and A is the original noise-free image. The background image, Ω is The defect region is defined by λ, which is a balancing parameter used to adjust the weights between the gradient domain term and the boundary consistency term.

[0065] This invention utilizes Poisson mixing to achieve natural blending in the gradient domain. By minimizing gradient differences and boundary consistency, it ensures that the edges and textures of the generated image remain consistent with the original image, while also ensuring smooth blending of the image with the background image at the boundaries.

[0066] Step 3.2. Wavelet domain texture synthesis.

[0067] The defect region of the defect image generated in step 2 is subjected to a third-level Haar wavelet decomposition. Low-frequency components from the background image are preserved to ensure the overall structure of the generated fused image is consistent with the background. Simultaneously, high-frequency texture details from the generated defect region are fused, resulting in a more refined and realistic texture. The formula for the third-level Haar wavelet decomposition is as follows: ; in, To merge images, for The low-frequency components of the background image represent the large-scale information of the image. This large-scale information includes the main outline, basic structure, and overall brightness distribution of the image, representing the most macroscopic structural information. For wavelet decomposition series, for High-frequency information in the horizontal direction of the defect area. for High-frequency information in the vertical direction of the defect area. for High-frequency information along the diagonal direction of the defect area.

[0068] The texture of the generated fused image is extracted using wavelet decomposition, preserving the low-frequency components of the background image to ensure the overall structure of the fused image matches the background. Adjusting or fusing the high-frequency components of the fused image further refines and enhances the realism of the generated region's texture. This helps maintain the texture continuity and detail of the fused image.

[0069] Step 3.3. Adaptive color migration.

[0070] Adaptive histogram matching is used to adjust the color distribution of the generated region, and the color of the fused image is adjusted by using the cumulative distribution function of the background image to eliminate the visual abruptness caused by lighting differences.

[0071] First, based on the fused image The color distribution of the generated image is compared with the color distribution of the background image using a cumulative distribution function; the colors of the generated image are adjusted by inverse mapping to match the color distribution of the background image, as shown in the following formula: ; The blended image after color adjustments. For defect images cumulative distribution function inverse function, Original noise-free image The cumulative distribution function of the background image.

[0072] This method maintains the naturalness of the generated area's color while eliminating visual abruptness caused by differences in lighting.

[0073] This invention employs a three-level fusion strategy to achieve seamless integration of defects with background lighting, texture, and color, ensuring visual consistency between the generated defects and the original image background. This results in a visually consistent fused image, as shown in the generated result. Figure 5 As shown.

[0074] Step 4. Build a segmentation network and pre-train it using an industrial defect dataset to obtain a pre-trained defect segmentation model.

[0075] The segmentation network consists of a hybrid encoder and a decoder. The hybrid encoder adopts the Mamba-Transformer-CBAM architecture, which combines convolution, Mamba modules, Transformer, and CBAM dual attention mechanisms.

[0076] This invention trains a segmentation network with the above-mentioned encoding and decoding structure on the publicly available MVTec industrial defect dataset, learns the feature representation and segmentation ability of general industrial defects, and obtains a pre-trained defect segmentation model.

[0077] Supervised pre-training based on the MVTec dataset includes two processes: 1. Data preparation and preprocessing: All images in the MVTec industrial anomaly detection dataset, which contains images of 15 types of industrial products and corresponding defect annotations, were uniformly resized to 512*512. Then, geometric transformations such as random cropping, rotation, and flipping were applied to enhance data diversity. The preprocessed images were then input into the pre-trained model architecture.

[0078] 2. Model Architecture Design: The pre-trained architecture is consistent with the formal training architecture, both employing an improved hybrid encoder-decoder architecture. The encoder integrates convolutional layers, Mamba modules, Transformer, and CBAM dual attention mechanisms.

[0079] like Figure 6 As shown, after local features are extracted from the input image data in the convolutional layer, the Mamba module is used to model long-range dependencies, the Transformer module captures global dependencies, and the CBAM dual attention module obtains the importance of feature maps in the channel and spatial dimensions, enhancing important features and suppressing unimportant features, thereby improving the model's feature representation ability and performance. The decoder adopts a multi-scale feature fusion design, combining skip connections and the CBAM attention mechanism to finally generate a probability map of defects.

[0080] This invention utilizes a hybrid encoder structure of Mamba-Transformer-CBAM to enable pre-trained models to capture multi-scale defect features, thereby providing high-quality initial parameters for transfer to rail data.

[0081] The defect segmentation model pre-trained in this step can provide high-quality initial parameters for fine-tuning in step 5.

[0082] Step 5. Supervised learning enhanced by adversarial training.

[0083] like Figure 7 As shown, an adversarial training framework is adopted, which consists of a generator and a discriminator. The pre-trained defect segmentation model is used as the generator in the generative adversarial network. The generator is fine-tuned using the rail defect dataset from step 1 to improve the boundary quality and realism of the segmentation mask generated by the generator. The discriminator constrains the realism of the generated mask.

[0084] An adversarial training strategy is employed to optimize the segmentation performance of rail surface defects. At the start of the adversarial training phase, the teacher model is obtained by deep copying the initial parameters of the student model (generator). Using the manually labeled real rail images from step 1 (i.e., the original noise-free images), the adversarial training framework is used to fine-tune the segmentation network model on the dataset. Adversarial game theory is utilized to improve the boundary quality and realism of the segmentation masks generated by the generator, resulting in an optimized student generator and a teacher model stabilized through EMA. The teacher model will be used in step 6 for the unlabeled data (i.e., the data generated in step 3). Generate pseudo tags.

[0085] The generator adopts the segmentation model architecture of step 4, fine-tunes the pre-trained model on the rail defect dataset, introduces an adversarial training mechanism to improve the quality of segmentation boundaries, and uses the principle of generative adversarial networks to make the segmentation results more consistent with the real defect distribution.

[0086] The process of supervised learning enhanced by adversarial training is as follows: 1. Data preparation.

[0087] The original, noise-free image with annotations is used as input for supervised learning.

[0088] 2. Adversarial training framework.

[0089] The generator uses the pre-trained defect segmentation model from step 4. The generator's input is an image of the rail surface, and its output is a defect segmentation probability map. The discriminator uses a convolutional neural network, progressively downsampling feature maps through 5 convolutional layers. Its input is an image-mask pair, and its output is the probability of a true segmentation. The discriminator determines whether the input image-mask pair comes from a ground-labeled dataset (labeled "true") or is a segmentation result generated by the generator (labeled "false"), and its output is a scalar probability value (representing the probability that the input pair is "true"). The generator and discriminator are trained alternately. First, the generator is fixed, and the discriminator is trained. A batch of real images and their manually labeled ground-labeled masks (labeled "true") and the segmentation masks generated by the generator from the same batch of images (labeled "false") are used to train the discriminator. The goal is to maximize the discriminator's accuracy in classifying real pairs (output close to 1) and minimize its accuracy in classifying generated pairs (output close to 0). Then, the discriminator is fixed, and the generator is trained. A batch of real images is input into the generator to obtain segmentation masks.

[0090] The goals of adversarial training are twofold: (1) to minimize the segmentation loss, making the generator's segmentation results as close as possible to the real annotations; and (2) to minimize the adversarial loss, ensuring that the generated mask and the original image can "deceive" the discriminator, i.e., maximizing the probability that the discriminator judges it as "real". The discriminator learns to distinguish between the expert-annotated perfect segmentation mask (real) and the generator-generated segmentation mask (fake). It strives to continuously improve its discrimination ability. The generator not only learns to generate accurate segmentation masks but also learns to generate masks that appear "indistinguishable from real" to the discriminator. It strives to deceive the discriminator. By introducing an adversarial game between the discriminator and the generator, the detail, structural coherence, and visual realism of the segmentation image can be improved.

[0091] 3. Teacher model generation.

[0092] At the start of the adversarial training phase, the teacher model is obtained by deep copying the initial parameters of the student model (generator). After each training round, the teacher model parameters are updated using an exponential moving average (EMA).

[0093] .in, Let be the teacher model parameters at step t. The parameters of the teacher model from the previous step (step t-1) are... Here are the parameters of the student model, and m is the smoothing coefficient.

[0094] This process establishes a stable parameter trajectory, enabling the teacher model to reach its optimal state during the integration training process. After completing adversarial training, the stable teacher model, updated via EMA, will be used for semi-supervised training guided by dynamic pseudo-labels in step 6. This design ensures that the teacher model, as a reliable knowledge carrier, can effectively transfer the composite knowledge from labeled and generated data.

[0095] Step 6. Semi-supervised training guided by dynamic pseudo-labels.

[0096] This invention generates pseudo-labels for fused images based on an EMA-stable teacher model, and identifies difficult samples through confidence screening, feature clustering, and dynamic reweighting. After applying data augmentation to the difficult samples, they are input into the student generator, i.e., the student model, for training. The teacher model smoothly updates its parameters through EMA, and finally obtains the optimal model weights.

[0097] By using a trained student model with optimal model weights, the surface defects of the rail can be segmented.

[0098] For the small number of real rail images with manual annotations in step 1 and the large number of unannotated fused images generated in step 3, dynamic pseudo-label generation, sample classification, and strong data augmentation are used to obtain the final optimized rail surface defect segmentation model, which improves the model's generalization ability in the field of rail defect detection and solves the core problem of insufficient labeled data.

[0099] Step 6.1. Dynamic pseudo-tag generation.

[0100] The obtained EMA-stable teacher model is used to generate pseudo-labels for the unlabeled data in the fused image generated in step 3.

[0101] Step 6.2. Sample Classification Strategy. Fine-grained classification is performed on pseudo-labeled samples to achieve differentiated training.

[0102] 1) Confidence level screening.

[0103] First, fuzzy samples are filtered out by calculating prediction entropy. Entropy measures the uncertainty of the probability distribution predicted by the model; the higher the entropy, the more uncertain the model is. The calculation formula is as follows: .

[0104] Where p is the class probability distribution of the model output.

[0105] The entropy of the prediction represents the model's uncertainty regarding the prediction outcome. (Choose...) Fuzzy samples, i.e., only entropy exists. Only samples within the specified range are considered difficult samples. The low confidence threshold This is the high confidence threshold.

[0106] 2) Feature space clustering.

[0107] Next, K-means clustering is performed on the deep semantic features of the fused image generated in step 3 using ResNet50. Samples close to the class boundaries are selected to supplement the confidence screening blind zone in order to capture geometric distribution features.

[0108] 3) Dynamic reweighting.

[0109] An adaptive weight is assigned to the loss for hard samples, adaptively adjusting the learning intensity. By adjusting the loss weights for hard samples, the model's learning of these samples is strengthened. (The text then repeats the "loss weights for hard samples" part.) The formula is as follows: .

[0110] in For the current batch precision, For moving average accuracy, This is the scaling factor, which controls the rate of change of the weights.

[0111] Through a three-step mechanism of confidence screening, feature space clustering, and dynamic reweighting, the model can focus on those difficult-to-classify samples, especially those at the classification boundary, which helps to improve the overall performance of the model.

[0112] Step 6.3. Strong data augmentation.

[0113] The set of difficult samples output from the three-stage screening provides input for strong data augmentation, forming an iterative optimization process of "screening-augmentation-learning". Strong data augmentation will perform color space transformation, geometric transformation, and spatial transformation on the screened difficult samples to improve sample diversity.

[0114] Step 6.4. Iterative optimization and result output.

[0115] The training process employs a two-stage optimization strategy. First, in the parameter update stage, to continuously improve the quality of pseudo-labels, the student model updates parameters through backpropagation, while the teacher model updates smoothly through step EMA.

[0116] To ensure the stability of the teacher model and to enable it to receive training parameters from the student model for effective updates, the smoothing coefficient m in the EMA is set to 0.999. The teacher model retains 99.9% of its historical parameters and 0.1% of the new parameters from the student model.

[0117] Then, during the evaluation and validation phase, the validation set mIoU is calculated at the end of each cycle, and the weights of the model with the best performance are saved.

[0118] Step 5 improves the quality of segmentation boundaries through adversarial training, and Step 6 utilizes the dynamic pseudo-label mechanism (confidence screening + clustering + reweighting) to efficiently utilize unlabeled generated data, solve the problem of scarce labels, and significantly enhance the generalization ability for small samples.

[0119] Using a trained student model (applying optimal model weights), surface defects on rails can be segmented, such as... Figure 8 The image shows the segmentation effect of track defects on the original image with a segmentation mask and the generated image without a segmentation mask. Figure 8 It can be seen that the present invention can achieve accurate segmentation of surface defects in both types of images.

[0120] Example 2 This embodiment 2 describes a rail surface defect detection system driven by a diffusion model, which is based on the same inventive concept as the rail surface defect detection method driven by a diffusion model in embodiment 1 above.

[0121] The rail surface defect detection system based on the diffusion model in this embodiment includes the following modules: The data acquisition and dynamic enhancement processing module is used to acquire the original noise-free image of the rail surface and generate the original defect mask through manual annotation. The original noise-free image and the corresponding original defect mask together form the rail defect dataset. For the original defect mask, dynamic mask generation is performed based on multi-level morphological transformation to obtain an enhanced defect mask; The defect generation module is designed to generate a noisy image with a fused preset defect shape by using an improved conditional diffusion model architecture for the original noise-free image and the enhanced defect mask. In the latent space of the image, the enhanced defect mask is used as a condition and the diffusion process is guided by dual-path control to generate a noisy image with a fused preset defect shape. After denoising, the defect image is obtained. The image seamless fusion processing module is used to fuse the defective region of the generated defective image into the background region of the original noise-free image using a three-level fusion strategy, so as to obtain a visually highly consistent fused image. The defect segmentation model building module is used to build a segmentation network that includes a hybrid encoder and a decoder; The hybrid encoder integrates convolution, Mamba modules, Transformer, and CBAM dual attention mechanisms; the decoder employs multi-scale feature fusion, combining skip connections and CBAM attention mechanisms to generate a defect segmentation probability map. The segmentation network is pre-trained on the MVTec dataset to obtain a pre-trained defect segmentation model. The supervised training module is used to employ an adversarial training framework, using the defect segmentation model as the generator in the generative adversarial network. The generator is fine-tuned using a rail defect dataset to obtain an optimized student generator and a teacher model that is stable through EMA. The semi-supervised training module is used to generate pseudo-labels for the fused images based on the EMA-stable teacher model. The hard samples are then identified through confidence screening, feature clustering, and dynamic reweighting. After data augmentation is applied to the hard samples, they are input into the student generator, i.e., the student model, for training. The teacher model updates its parameters smoothly through EMA, and finally obtains the optimal model weights. By using a trained student model with optimal model weights, the surface defects of the rail can be segmented.

[0122] It should be noted that in the diffusion model-driven rail surface defect detection system of this embodiment 2, the implementation process of the functions and roles of each functional module is detailed in the implementation process of the corresponding steps of the method in embodiment 1 above, and will not be repeated here.

[0123] Example 3 This embodiment 3 describes a computer device. The computer device includes a memory and one or more processors. Executable code is stored in the memory. When the processor executes the executable code, it implements the steps of the diffusion model-driven rail surface defect detection method described in embodiment 1 above.

[0124] In this embodiment, the computer device can be any device or apparatus with data processing capabilities, and will not be described in detail here.

[0125] Example 4 This embodiment 4 describes a computer-readable storage medium storing a program that, when executed by a processor, is used to implement the steps of the diffusion model-driven rail surface defect detection method in embodiment 1 above.

[0126] The computer-readable storage medium can be an internal storage unit of any device or apparatus with data processing capabilities, such as a hard disk or memory, or an external storage device of any device with data processing capabilities, such as a plug-in hard disk, smart media card (SMC), SD card, flash card, etc.

[0127] Of course, the above description is only a preferred embodiment of the present invention. The present invention is not limited to the above-described embodiments. It should be noted that any equivalent substitutions or obvious modifications made by those skilled in the art under the guidance of this specification fall within the scope of this specification and should be protected by the present invention.

Claims

1. A method for detecting rail surface defects based on a diffusion model, characterized in that, Includes the following steps: Step 1. Acquire the original noise-free image of the rail surface and generate the original defect mask through manual annotation. The original noise-free image and the corresponding original defect mask together form the rail defect dataset. For the original defect mask, dynamic mask generation is performed based on multi-level morphological transformation to obtain an enhanced defect mask; Step 2. For the original noise-free image and the enhanced defect mask, an improved conditional diffusion model architecture is adopted. In the latent space of the image, the enhanced defect mask is used as a condition, and the diffusion process is guided by dual-path control to generate a noisy image with a fused preset defect shape. After denoising, the defect image is obtained. In step 2, a latent variable diffusion process is defined, and the cumulative noise coefficient is set by using a noise addition formula and a cosine scheduler to retain more signal in the early stage of training in order to improve convergence stability and obtain a noisy image. The diffusion model adopts the UNet architecture, which has a symmetrical structure and includes an encoder, intermediate blocks, and a decoder. The encoder part includes multiple encoding blocks, and the decoder part includes multiple decoding blocks; Dual-channel control modules are set in some coding blocks and intermediate blocks of UNet. The input of the dual-channel control modules are the output features of the edge control module after downsampling and the text embedding features output by the material control module. The edge control module extracts edge features from the original noise-free image using the Canny operator, outputs a feature map after processing through multiple residual blocks, and then downsamples it to serve as input to the UNet partial coding block and intermediate block of the diffusion model, respectively. The material control module is used to receive material text input and obtain text embedding features through a text encoder; The dual-path control module performs cross-attention on the output features of the downsampled edge control module and the text embedding features output by the material control module, and outputs a feature map with edge information and material information to guide the diffusion process. UNet predicts and removes noise at multiple time steps to generate high-quality defect images; By jointly optimizing the noise prediction loss, perception loss, and contrast loss, the diffusion model generates images that are both realistic and effectively highlight defect areas, while maintaining consistency in the noise distribution of the latent space. Step 3. For the defect area of ​​the generated defect image and the background area of ​​the original noise-free image, a three-level fusion strategy is used to fuse the defect area into the background area to obtain a visually highly consistent fused image; Step 3 specifically involves: Step 3.

1. Using the background of the original noise-free image as a reference, based on gradient domain Poisson fusion, the defect region of the defect image generated in Step 2 is fused to the background to minimize gradient differences and boundary inconsistencies. Step 3.

2. Wavelet domain texture synthesis: Perform three-level Haar wavelet decomposition on the defect region of the defect image generated in Step 2, retain the low-frequency part of the background image to ensure that the overall structure of the generated fused image is consistent with the background, and at the same time fuse the high-frequency texture details of the generated defect region to make the texture of the generated region more refined and realistic. Step 3.

3. Adaptive color migration: Adjust the color distribution of the generated region using adaptive histogram matching. Adjust the color of the fused image by using the cumulative distribution function of the background image to eliminate the visual abruptness caused by lighting differences. Step 4. Build a segmentation network and pre-train it using an industrial defect dataset to obtain a pre-trained defect segmentation model; Step 5. Using an adversarial training framework, the pre-trained defect segmentation model is used as the generator in the generative adversarial network. The generator is fine-tuned using the rail defect dataset to obtain an optimized student generator and a teacher model that is stable through EMA. Step 6. Based on the stable teacher model of EMA, pseudo-labels are generated for the fused image, and difficult samples are identified through confidence screening, feature clustering, and dynamic reweighting. After data augmentation is applied to the difficult samples, they are input into the student generator, i.e., the student model, for training. The teacher model updates its parameters smoothly through EMA, and the optimal model weights are obtained through training. By using a trained student model with optimal model weights, the surface defects of the rail can be segmented.

2. The rail surface defect detection method based on diffusion model driven according to claim 1, characterized in that, In step 1, the process of dynamic mask generation based on multi-level morphological transformation is as follows: The original defect mask is subjected to random morphological operations, multi-scale Gaussian noise injection, and random rotation transformation in sequence to simulate the morphological variability and irregular boundaries of real defects, and dynamic mask generation is performed. In a random morphological operation sequence, each operation kernel randomly selects an elliptical or rectangular structure, and the kernel size is randomly generated; Random noise is injected across multiple scales to simulate the irregular boundaries of real defects; The spatial diversity is enhanced by applying random rotation transformation, with the rotation angle randomly selected between -45° and 45°.

3. The rail surface defect detection method based on diffusion model driven according to claim 1, characterized in that, In step 2, the total loss function of the diffusion model The calculation formula is: ； ； ； ； in, Predict loss for noise; To mitigate perceptual loss, the generated image is made perceptually similar to the real image. This is for contrast loss, used to enhance the contrast between the defective area and the background; , , These are the weighting coefficients for each loss term; Noise predicted by UNet; These are the parameters of the network; The image contains noise at time step t; t is a random time step. c represents real noise; 'c' represents conditional input. Indicates the first The number of channels in the layer feature map; Indicates the first The height of the layer feature map; Indicates the first The width of the layer feature map; Indicates the first Layer features; Defect images generated for the diffusion model; The original, noise-free image; The feature vector of the defect region; Background features; τ is the feature vector of the k-th region, used to calculate the contrast between the background and the defect; τ is the temperature coefficient, which controls the smoothness of the probability distribution. This indicates the total number of regions.

4. The rail surface defect detection method based on diffusion model driven according to claim 1, characterized in that, In step 3.1, the formula for minimizing gradient difference and boundary inconsistency is as follows: ； in, The original noise-free image, B is the defect image generated in step 2, and A is the original noise-free image. Background image, Ω is The defect region, where λ is a balance parameter used to adjust the weights between the gradient domain term and the boundary consistency term; In step 3.2, the formula for performing the three-level Haar wavelet decomposition is as follows: ； in, To merge images, for The low-frequency components of the background image; For wavelet decomposition series, for High-frequency information in the horizontal direction of the defect area. for High-frequency information in the vertical direction of the defect area. for High-frequency information along the diagonal direction of the defect area; In step 3.3, based on the fused image The color distribution of the generated image is compared with the color distribution of the background image using a cumulative distribution function; the colors of the generated image are adjusted by inverse mapping to match the color distribution of the background image, as shown in the following formula: ； in, The blended image after color adjustments. For defect images cumulative distribution function inverse function, Original noise-free image The cumulative distribution function of the background image.

5. The rail surface defect detection method based on diffusion model driven according to claim 1, characterized in that, In step 4, the segmentation network includes a hybrid encoder and a decoder; The hybrid encoder integrates convolution, Mamba modules, Transformer, and CBAM dual attention mechanisms; the decoder employs multi-scale feature fusion, combined with skip connections and CBAM attention mechanisms, to generate a defect segmentation probability map. After the image input to the hybrid encoder extracts local features in the convolutional layer, the Mamba module is used to model long-range dependencies, the Transformer module captures global dependencies, and the CBAM dual attention module obtains the importance of feature maps in the channel and spatial dimensions. The segmentation network is pre-trained on the MVTec dataset to obtain a pre-trained defect segmentation model.

6. The rail surface defect detection method based on diffusion model driven according to claim 1, characterized in that, In step 5, the generator uses the defect segmentation model pre-trained in step 4. The input of the generator is the original noise-free image of the rail surface, and the output of the generator is the defect segmentation probability map. The discriminator uses a convolutional neural network. Through 5 layers of convolutional network, the feature map is downsampled step by step. The input of the discriminator is the image and mask pair, and the output is the probability of the segmentation's authenticity. The discriminator is used to determine whether the input image and mask pair comes from the real labeled dataset, i.e., the rail defect dataset, or is the segmentation result generated by the generator. The output is a scalar probability value, which represents the probability that the input image and mask pair is real. Step 5 employs an adversarial training strategy to optimize the segmentation performance of rail surface defects. At the beginning of the adversarial training phase, the teacher model obtains its initial parameters by deeply copying the student model, i.e., the generator. For the real rail images with manual annotations from Step 1, i.e., the original noise-free images, an adversarial training framework is used to fine-tune the segmentation network model on the dataset. The adversarial game is used to improve the boundary quality and realism of the segmentation mask generated by the generator, resulting in an optimized student generator and a teacher model stabilized through EMA. The teacher model will be used in Step 6 to generate pseudo-labels for the unlabeled data, i.e., the fused images.

7. The rail surface defect detection method based on diffusion model driven according to claim 1, characterized in that, Step 6 specifically involves: Step 6.

1. Generate pseudo-labels for the fused image from Step 3 based on the EMA-stable teacher model obtained in Step 5; Step 6.

2. For pseudo-labeled samples, use confidence screening, feature space clustering and dynamic reweighting mechanism to perform fine classification, so that the model can focus on those difficult-to-classify samples and obtain the difficult sample set; Step 6.

3. Perform color space transformation, geometric transformation, and spatial transformation on the selected difficult sample set to improve sample diversity; Step 6.

4. Iterative optimization and result output: The student model updates its parameters through backpropagation, and the teacher model updates smoothly through EMA. After each cycle, the validation set mIoU is calculated, and the weights of the model with the best performance are saved.

8. The rail surface defect detection method based on diffusion model driven according to claim 7, characterized in that, In step 6.4, the training process adopts a two-stage optimization strategy. First, in the parameter update stage, in order to continuously improve the quality of pseudo-labels, the student model updates the parameters through backpropagation, and the teacher model updates them smoothly through step EMA. To ensure the stability of the teacher model and to effectively update the training parameters of the student model, the smoothing coefficient m in the EMA is set to 0.999, and the teacher model retains 99.9% of the historical parameters and 0.1% of the new parameters of the student model. Then, during the evaluation and validation phase, the validation set mIoU is calculated at the end of each cycle, and the weights of the model with the best performance are saved.

Images