Semi-supervised medical image segmentation method based on single-step cold diffusion model

Abstract

The invention discloses a semi-supervised medical image segmentation method based on a single-step cold diffusion model, and solves the problems of dependence on a large amount of annotation data, pseudo label error accumulation and low diffusion model efficiency in the prior art. The method constructs a "forward deterministic degradation-reverse progressive recovery-single-step distillation" architecture: standardizing a medical image, binarizing a mask and dividing data; a cold diffusion model is constructed based on U-Net, and Gaussian noise is replaced by 15-step structured shielding; designing a progressive degeneration dividing and treating strategy according to the geometric center of the focus and the nonlinear curve; a multi-step cold diffusion model is used as a teacher generation pseudo label, a single-step student model is distilled, and segmentation and consistency loss optimization are carried out. The method is high in precision under the condition of few labels, adapts to MRI, CT and ultrasonic multi-mode and brain, kidney and breast scenes, and is close to real-time reasoning and high in robustness.

Description

Technical Field

[0001] This invention relates to the field of medical image processing technology, and designs a semi-supervised medical image segmentation method based on a single-step cold diffusion model, belonging to the field of image segmentation. Background Technology

[0002] Medical image segmentation is a core technical support for clinical diagnosis, treatment planning, and prognostic assessment. Its goal is to accurately separate lesion regions from normal tissue, providing quantitative evidence for medical decisions. Traditional medical image segmentation methods mainly rely on fully supervised learning. Among them, UNet and its variants have become benchmark models in this field due to their efficient encoder-decoder structure and skip connection design, and are widely used in image segmentation tasks such as CT and MRI. However, fully supervised learning heavily depends on large-scale, accurately labeled data. Medical image annotation requires professional physicians to spend a lot of time and effort, and individual differences in lesion morphology and pathological diversity further increase the annotation difficulty, resulting in a scarcity of high-quality labeled datasets. This greatly limits the generalization ability and clinical applicability of traditional UNet models.

[0003] To alleviate the dependence on labeled data, semi-supervised medical image segmentation methods have emerged. Their core idea is to train the model using a small amount of labeled data combined with a large amount of unlabeled data. Existing semi-supervised methods often utilize unlabeled data through pseudo-labeling strategies. That is, the model first generates predicted labels for the unlabeled data, and then uses these as "pseudo-labels" to train alongside the real labeled data. However, such methods generally suffer from inaccurate pseudo-labels, and erroneous pseudo-labels accumulate and propagate during iterative training, leading to a decline in model performance.

[0004] In recent years, diffusion models have been introduced into the field of medical image segmentation due to their powerful generative modeling capabilities. Unlike traditional UNets that directly learn the mapping from input to output, diffusion models implicitly learn the data distribution through an iterative process of "adding noise and denoising," enabling the generation of diverse segmentation results and demonstrating significant advantages in optimizing segmentation edge details. This provides a new path to solve the pseudo-label error problem in semi-supervised segmentation. However, the design of classic diffusion models has two major drawbacks: First, the forward process relies on multi-step random perturbations of Gaussian noise, requiring dozens or even hundreds of sampling iterations to ensure denoising quality, resulting in extremely low training and inference efficiency, making it difficult to meet the real-time requirements of clinical practice. Second, the globally uniform noise addition method does not consider the extreme imbalance between lesions (foreground) and normal tissue (background) in medical images, consuming a large amount of computational resources in background regions that contribute very little to the segmentation task, resulting in serious computational redundancy.

[0005] Furthermore, existing semi-supervised segmentation methods based on diffusion models fail to fully utilize the structured prior knowledge of medical images, and still suffer from insufficient pseudo-label accuracy and a difficulty in balancing computational efficiency and segmentation performance. Therefore, there is an urgent need for a semi-supervised medical image segmentation technique that combines high accuracy, high efficiency, and strong robustness to overcome the bottlenecks of existing methods and adapt to practical clinical applications. Summary of the Invention

[0006] To address the shortcomings of existing technologies, this invention provides a semi-supervised medical image segmentation method based on a single-step cold diffusion model. This invention aims to solve three core problems existing in current semi-supervised medical image segmentation techniques: 1) Traditional UNet models rely on large-scale labeled data and have poor generalization ability in scenarios with few labeled data; 2) The pseudo-labels of existing semi-supervised methods are inaccurate and prone to cumulative errors; 3) The "noise addition-denoising" paradigm of traditional diffusion models leads to low computational efficiency, serious computational redundancy, and difficulty in balancing segmentation accuracy and inference speed.

[0007] To achieve the above objectives, this invention discloses an efficient semi-supervised medical image segmentation framework based on a single-step cold diffusion model. Its core is the reconstruction of the traditional diffusion paradigm, achieving a balance between high accuracy and high efficiency through a deterministic degradation-recovery mechanism and a single-step distillation strategy. The specific technical solution is as follows:

[0008] Overall framework design as Figure 1 As shown, based on the cold diffusion model, a three-level architecture of "forward deterministic degradation - inverse progressive recovery - single-step distillation" is constructed. The forward process gradually obscures the original segmentation mask through a structured degradation operator, the inverse process reconstructs the target mask through a parameterized recovery network, and finally, a lightweight model for single-step inference is obtained through distillation learning.

[0009] The cold diffusion model's core module abandons the Gaussian noise perturbation of classic diffusion models and employs a generalized deterministic forward degradation operator sequence. This sequence progressively processes the original data through a series of structured degradation transformations, forming a degraded data sequence. The training objective is to learn an inverse recovery function that directly predicts the original data based on the degraded data and the corresponding degradation steps. Model training is completed by minimizing the reconstruction error between the recovered and original data, thus overcoming probabilistic sampling constraints and laying the foundation for efficient inference.

[0010] Progressive divide-and-conquer degradation strategy: Design a goal-guided nonlinear degradation operator, the specific steps of which include:

[0011] First, the foreground target region in the original segmentation mask is determined, and the geometric center of this region is calculated, which is used as the origin of the degradation process. A normalized distance field is constructed to quantify the positional importance of each foreground pixel relative to the geometric center, while background pixels are not considered in the degradation process. A nonlinear progress function is designed to control the degradation rhythm. This function has a low growth rate in the early stage of degradation, which fully models the fine structure of the target contour, and the growth rate increases significantly in the later stage, which accelerates the masking of the main target region.

[0012] Based on the definition of a dynamic distance threshold for each step using a nonlinear progress function, an occlusion mask for the corresponding step is generated. By performing element-wise operations between the occlusion mask and the original mask and the original image, the degraded mask and the degraded image are obtained.

[0013] Single-step distillation learning paradigm: A multi-step cold diffusion model is used as the teacher model, leveraging its efficiency to generate high-precision pseudo-labels for a large amount of unlabeled data; a lightweight single-step student model is trained, with two learning objectives: first, to mimic the final segmentation output of the teacher model on labeled data; and second, to maintain consistency with the pseudo-labels generated by the teacher model on unlabeled data, optimizing the student model parameters through a weighted sum of segmentation loss and consistency loss. Attached Figure Description

[0014] Figure 1 Schematic diagram of a semi-supervised medical image segmentation model based on single-step cold diffusion

[0015] Figure 2 Visualization results of the BraTS2018 comparative trial

[0016] Figure 3 Visualization results of KiTS19 comparative test

[0017] Figure 4 Visualization results of the BUSI comparative test Detailed Implementation

[0018] S1: Data preprocessing;

[0019] This implementation method is applicable to three typical medical image datasets (BraTS2018, KiTS19, BUSI), and the preprocessing steps are uniformly as follows:

[0020] S11. Image Standardization: Normalize the pixel values ​​of the input medical image to the [0,1] interval, using the following formula:

[0021]

[0022] Where I(i,j) is the original image pixel value, and min(I) and max(I) are the global minimum and maximum pixel values ​​of the image, respectively;

[0023] S12. Mask Binarization: The segmentation mask M0 is processed into a binary matrix, with the foreground lesion region labeled as 1 and the background region labeled as 0, i.e., M0∈{0,1}. H×W (H and W are the image height and width, respectively);

[0024] S13. Data partitioning: According to the annotation ratio (10% of the annotation data D) L 90% of the data is unlabeled. U or 20% D L 80% D U ) Divide the dataset, D L Includes image-mask pairs (I,M), D U Image I only.

[0025] S2: Implementation of the cold diffusion model;

[0026] The forward process is a deterministic sequence x t =D t (x0), t = 1, 2, ..., T, where x0 is the original data (image I or mask M0, D) t For degradation transformation (structured occlusion in this invention), T is the degradation step number (value 15); Inverse recovery operator (U-Net network): Recovery operator R θ (·) The U-Net architecture is adopted, with specific structural parameters: Encoder: 4 downsampling modules, each containing 2 residual blocks, the residual blocks integrating group normalization and the Swish activation function; downsampling uses convolution operations with a stride of 2, and the number of channels gradually doubles from 64 to 512; Decoder: 4 upsampling modules, each containing 2 residual blocks, upsampling uses transposed convolutions (stride of 2), and the number of channels gradually halves from 512 to 64; Attention mechanism: self-attention modules are introduced in layers 3 and 4 of the encoder and layers 2 and 3 of the decoder to capture long-range dependencies; Temporal step embedding: sinusoidal positional encoding is used to map time step t to a feature vector, the formula is: PE(t,2k)=sin(t / 10000 2k / d ),PE(t,2k+1)=cos(t / 10000 2k / d Where d = 256 is the embedding dimension, and after mapping, it is injected into each residual block through a linear layer; Input / Output: Input is image I and degradation mask O t The channels are stitched together (total number of channels = number of image channels + 1), and the output is the restored segmentation mask.

[0027] S3, Implementation of the progressive divide-and-conquer degeneracy strategy;

[0028] This strategy guides occlusion based on the target's geometric features. The specific steps are as follows:

[0029] S31. Calculation of Foreground Region and Geometric Center: Define the foreground region Ω = {(i,j) | M0(i,j) = 1}, and calculate its geometric center.

[0030]

[0031] Where |Ω| represents the total number of foreground pixels;

[0032] S32. Normalized distance field construction: For each foreground pixel (i,j)∈Ω, calculate its distance relative to the center. Normalized distance:

[0033]

[0034] Background pixels The distance field is undefined and is not processed during degradation.

[0035] S33. Nonlinear Scheduling Curve Design: Define a linear schedule s = t / (T-1), t∈[0,T-1], and map it to a nonlinear schedule using an exponential function: f γ (s)=s γ γ = 2.0 dynamic distance threshold τ t for:

[0036]

[0037] S34. Occlusion Mask and Degradation Data Generation: Step t Occlusion Mask O t Generated via indicator function: O t (i,j)=I(D(i,j)≤τ t )·I((i,j)∈Ω) where I(·) is an indicator function (outputs 1 if the condition is met, otherwise 0); degenerate image I t With degradation mask M t For: M t =M0⊙O t ,I t =I⊙O t +μ·(1-O t ) where ⊙ represents element-wise multiplication, and μ = 0 (background pixel value).

[0038] S4, achieved through single-step distillation;

[0039] The distillation process compresses the 15-step teacher model into a single-step student model. The specific process is as follows:

[0040] S41. Teacher Model Training:

[0041] During teacher model training, the objective is to minimize the combination of Dice loss and cross-entropy loss between the recovered mask and the original mask. The loss formula is as follows: Dice loss:

[0042]

[0043] Cross-entropy loss:

[0044]

[0045] The optimizer used is AdamW, with a learning rate η = 2 × 10⁻⁶. -5 Weight decay λ = 1 × 10 -4 The training parameters are set to batch size of 4, gradient accumulation step size of 2, and number of training iterations of 10. 5 After training, the teacher model will be able to apply D. U Perform 15-step restoration on unlabeled images to output high-precision pseudo-labels.

[0046] S42, Student Model Training:

[0047] The student model architecture is completely identical to the teacher model's U-Net architecture, except that iterative sampling modules are removed, while retaining the single forward propagation path. Its total loss function integrates the segmentation loss on labeled data and the teacher-student consistency loss on unlabeled data, as shown in the formula: Where F φ (·) represents the student model, φ represents the student model parameters, the consistency loss uses the same Dice + cross-entropy combination as the segmentation loss, the training parameters are the same as the teacher model, and the number of iterations is 5 × 10. 4 (Based on fine-tuning of the pre-trained teacher model).

[0048] S5. Model Evaluation;

[0049] S51. Dataset Selection:

[0050] We selected three major public datasets: BraTS2018 (brain MRI), KiTS19 (kidney CT), and BUSI (breast ultrasound), covering 2D / 3D dimensions and multimodal images. The training and test sets were divided into training and test sets according to a reasonable ratio. After standardization, cropping and other preprocessing, the effectiveness and generalization of the model were fully verified.

[0051] S52. Evaluation Indicators:

[0052] The Dice Similarity Coefficient (DSC), Jaccard Index, Hausdorff Distance (HD), and Average Symmetric Surface Distance (ASD) are used as the core evaluation metrics. The calculation formulas are as follows:

[0053] DSC: Quantitative prediction of overlap with the true label, the formula is as follows.

[0054]

[0055] A and B are the binary mask sets of the predicted segment and the real label, respectively. |A∩B| is the size of the intersection of the two. |A| and |B| are the cardinality of their respective sets, with a value range of 0-1. The higher the value, the better the overlap.

[0056] Jaccard(IoU): Measures pure overlap, the formula is as follows

[0057]

[0058] The variables have the same meaning as DSC, but supplement DSC, making them more sensitive to boundary deviations.

[0059] HD: Hausdorff distance for outlier removal, HD(A,B)=HD(A,B)=max(h(A,B),h(B,A)), where

[0060]

[0061] A and B are sets of binary masks, h(A,B) is the maximum value of the shortest distance from all points in A to B, d(a,b) is the Euclidean distance between points a and b, and 95HD is the 95th percentile of HD. The smaller the value, the more accurate the boundary.

[0062] ASD: Average Deviation of Quantization Boundaries, the formula is...

[0063]

[0064] S_A and S_B are the surface point sets of A and B, respectively, |S_A| and |S_B| are the number of surface points, and d is the Euclidean distance. This smooths the evaluation of geometric quality and compensates for the shortcomings of HD.

[0065] S53, Ablation Test:

[0066] To verify the effectiveness of the core module of this invention, ablation experiments were conducted on the BUSI and BraTS2018 datasets (10% labeled + 90% unlabeled) to compare the effects of different degradation mechanisms and single-step distillation (SSD). The results are shown in Table 1. The comparison mechanisms include: 1) GGND: traditional global Gaussian noise degradation; 2) GGMD: global Gaussian mask degradation under the cold diffusion framework; 3) PDCD: the structured local degradation based on the lesion geometric center and nonlinear progression curve proposed in this invention.

[0067] Experiments show that GGND's Dice is below 50%, indicating severely inadequate performance. Replacing GGMD with the deterministic masking operation GGMD improves the Dice on the BUSI dataset to 70.98, confirming the applicability of the cold diffusion paradigm. PDCD further improves the Dice on the BUSI dataset to 72.82, while reducing 95HD and ASD by 4.23 and 2.91 respectively, verifying the superiority of structured degradation in lesion morphology modeling. Introducing SSD improves the Dice on the BraTS2018 dataset from 75.34 to 83.47, and reduces ASD to 3.31, achieving single-step compression of multi-step inference capabilities while balancing accuracy and inference efficiency.

[0068] Table 1 shows the ablation experiment results using 10% labeled data on the BUSI and BraTS2018 datasets.

[0069]

[0070] S54. Comparative Test:

[0071] This invention was experimentally validated on three commonly used medical image datasets, and the results are shown in Tables 2, 3, and 4, respectively. With a 10% annotation setting on the BraTS2018 dataset, this invention significantly outperforms the CGDM model: Dice reaches 83.47%, Jaccard reaches 72.54%, 95HD is reduced to 10.36, and ASD is reduced to 3.31, requiring only 15 sampling steps, balancing accuracy and inference efficiency. This is primarily due to the progressive divide-and-conquer degradation strategy that alleviates edge blurring and provides accurate pseudo-labels. With 10% / 20% annotation ratios on the KiTS19 and BUSI datasets, this invention outperforms all comparable methods in Dice, Jaccard, 95HD, and ASD metrics. It is adaptable to multimodal MRI, CT, and ultrasound imaging, as well as segmentation of multiple brain, kidney, and breast sites, demonstrating outstanding robustness and generalization.

[0072] To visually demonstrate the differences in segmentation performance, Figure 2-4Visual comparisons of the present invention with mainstream semi-supervised segmentation methods on the BraTS2018, KiTS19, and BUSI datasets are presented. It is evident that the segmentation mask of the present invention is closest to the ground truth annotation, with accurate and complete boundaries, and stronger robustness against noise. In the complex multi-lesion scenario of KiTS19, the present invention possesses a unique advantage: it is the only method capable of completely segmenting all discrete lesions, while other comparative methods suffer from incomplete segmentation or over-fusion. This confirms the effectiveness of the present invention's progressive divide-and-conquer degradation strategy, which, through deterministic iteration from the main body to the edge, prioritizes the reconstruction of the lesion core before refining the edges, ensuring the completeness and accuracy of segmentation in multi-target, morphologically variable scenarios.

[0073] Table 2 shows the segmentation performance on the BraTS2018 dataset.

[0074]

[0075]

[0076] Table 3 shows the segmentation performance on the KiTS19 dataset.

[0077]

[0078]

[0079] Table 4 shows the segmentation performance on the BUSI dataset.

[0080]

[0081]

[0082] S55. Comparative Test:

[0083] To evaluate the inference efficiency of this invention, it was compared with other medical image segmentation methods based on diffusion models, and the results are shown in Table 5.

[0084] Table 5 Efficiency Comparison Results

[0085]

Claims

1. A semi-supervised medical image segmentation method based on a single-step cold diffusion model, characterized in that, Includes the following steps: S1. Data Preprocessing: The medical image dataset is standardized and masked for binarization, and then divided into labeled datasets D according to a preset annotation ratio. L and unlabeled dataset D U The D L Containing image-mask pairs (I,M), the D U Image I only; S2. Cold Diffusion Model Construction and Training: Constructing the Recovery Operator R Based on the U-Net Architecture θ (·), defining a deterministic forward degenerate sequence x t =D t (x0), (t=1,2,...,T), where x0 is the original data, D t For structured occlusion degradation transformation, T is the degradation step number; the restoration operator is trained by minimizing the reconstruction error between the restored mask and the original mask to achieve inverse progressive restoration; S3, Implementation of the progressive divide-and-conquer degeneracy strategy: S31. Determine the foreground region Ω of the original segmentation mask M0, and calculate the geometric center of Ω. S32. Construct the normalized distance field D(i,j) of the Ω to quantify the positional importance of the foreground pixel relative to the geometric center; S33. Design a nonlinear progress function f γ (s)=s γ Based on the function, a dynamic distance threshold τ is defined. t ; S34, Generate occlusion mask O t Through the O t The degradation mask M is obtained by element-wise operation with the original mask and the original image. t and degraded image I t ; S4, Single-step distillation training: S41. Using the trained multi-step cold diffusion model as the teacher model, for D U Generating pseudo-labels for unlabeled images in the dataset; S42. Construct a single-step student model F φ (·), with the weighted sum of partitioning loss and consistency loss as the optimization objective, using D L Training student models with pseudo-labels; S5. Segmentation Inference: Input the medical image to be segmented into the trained student model, and output the segmentation result after a single forward propagation.

2. The method according to claim 1, characterized in that, In step S1, image standardization normalizes pixel values ​​to the [0,1] range; the annotation ratio is 10% annotated data + 90% unannotated data or 20% annotated data + 80% unannotated data.

3. The method according to claim 1, characterized in that, In step S2, the U-Net architecture encoder contains four downsampling modules, each containing two residual blocks with ensemble group normalization and Swish activation functions. Downsampling uses convolution operations with a stride of 2, and the number of channels gradually doubles from 64 to 512. The decoder contains four upsampling modules, each containing two residual blocks. Upsampling uses transposed convolutions, and the number of channels gradually halves from 512 to 64. Self-attention modules are introduced in the 3rd and 4th layers of the encoder and the 2nd and 3rd layers of the decoder. Time step t is mapped to a feature vector through sinusoidal position encoding, with the formula PE(t,2k)=sin(t / 10000). ( 2k / d)), PE(t,2k+1)=cos(t / 10000 ( 2k / d), where d = 256 is the embedding dimension, is mapped and injected into each residual block through a linear layer; the input to the recovery operator is the image I and the degradation mask O. t The channels are spliced ​​together, and the output is the restored segmentation mask M0.

4. The method according to claim 1, characterized in that, In step S31, the geometric center The calculation formula is |Ω| represents the total number of foreground pixels; in step S32, the normalized distance field is: Background pixels In step S33, the linear schedule s = t / (T-1), t ∈ [0, T-1], is mapped to a nonlinear schedule through an exponential function: f γ (s)=s γ γ = 2.0 dynamic distance threshold τ t for: In step S34, the masking code O t (i,j)=I(D(i,j)≤τ t )·I((i,j)∈Ω), where I(·) is the indicator function, and the degenerate image I t =I⊙O t +μ·(1-O t M t =M0⊙O t ⊙ represents element-wise multiplication, μ = 0.

5. The method according to claim 1, characterized in that, In step S41, the training loss of the teacher model is a combination of Dice loss and cross-entropy loss. In step S42, the student model structure is completely consistent with the U-Net architecture of the teacher model, except that the iterative sampling-related modules are removed, and the single forward propagation path is retained. Its total loss function integrates the segmentation loss on labeled data and the teacher-student consistency loss on unlabeled data.

Images