An efficient medical image segmentation method based on consistency learning

Abstract

This invention discloses an efficient medical image segmentation method based on consistency learning. It addresses the problems of slow inference speed and numerous sampling steps in existing diffusion-based medical image segmentation methods. This invention introduces an independent UNet structure as a supervised signal generator to extract multi-scale contextual features from the input medical image. These features are then injected into the decoding path of the main consistency model through a channel attention mechanism, enhancing the model's ability to perceive fine-grained anatomical structures (especially small lesions). In addition to the original consistency loss (ensuring consistency of outputs at different time steps), this invention introduces a direct segmentation supervision loss to form a joint optimization objective. This effectively guides the model to learn accurate segmentation boundaries and region semantics while maintaining single-step inference efficiency. This invention belongs to the field of medical image segmentation technology.

Description

Technical Field

[0001] This invention relates to an efficient medical image segmentation method, belonging to the field of medical image segmentation technology. Background Technology

[0002] Diffusion models have shown great potential in medical image segmentation tasks, generating images with strong noise resistance and smoothness. However, mainstream diffusion models require a large number of resampling steps (usually hundreds or even thousands) during training and prediction, resulting in low computational efficiency and lengthy inference time, which severely restricts their deployment in real-time clinical applications.

[0003] Recently, the Consistency Model, as an independent generative network, has effectively addressed this bottleneck issue of diffusion models. Unlike traditional diffusion models, which rely on a stepwise denoising process, the Consistency Model learns a generative function that satisfies a "consistency path" by forcing the model's output to remain consistent between any two adjacent time steps. This method significantly reduces the number of sampling steps to one (i.e., single-step inference), achieving generative quality comparable to diffusion models while greatly accelerating the training and prediction processes, demonstrating a superior efficiency-performance balance in image generation tasks.

[0004] Although the consistency model performs well in general image generation tasks, it is not directly applicable to image segmentation tasks. This is because: (1) segmentation tasks require models to accurately recover pixel-level semantic structures, rather than generating realistic images unconditionally or conditionally; (2) the original framework of the consistency model lacks an effective mechanism for utilizing the semantic content of the input image, making it difficult to establish an accurate mapping from medical images to their corresponding anatomical structure masks; and (3) in the field of medical imaging, target structures are often scale-variable, have blurred boundaries, and exhibit inter-class imbalance, posing a greater challenge to models that rely solely on consistency constraints. Therefore, the applicability of the consistency model in medical image segmentation has not yet been explored, and there is a lack of targeted architectural design and optimization strategies. Summary of the Invention

[0005] To address the issues of slow inference speed and numerous sampling steps in existing diffusion-based medical image segmentation methods, this invention proposes an efficient medical image segmentation method based on consistency learning.

[0006] The technical solution adopted by the present invention to solve the above problems is as follows: The steps of the present invention include:

[0007] Step 1: Construct a medical image dataset, standardize the original images and their corresponding segmentation masks, and divide them into training and test sets; Step 2: Select real segmentation masks from the training set as target samples, and construct noisy inputs in the continuous time domain using Gaussian noise; Step 3: Design a multi-scale supervised signal generation module to extract multi-scale contextual features from the original medical image and adaptively fuse them into the decoding path of the consistency master model through a channel attention mechanism to generate preliminary segmentation predictions; Step 4: Construct a consistent training framework, including a main model M and a target model TM, where the parameters of TM are updated by the exponential moving average of M; calculate the parameters of the main model M at time step... The output and the target model TM at adjacent time steps The consistency loss between the outputs is combined with the direct segmentation supervision loss to form a joint optimization objective; Step 5: During the testing phase, only one forward propagation is required: input the pure noise, multi-scale supervision signal and initial time step into the main model M to train, and directly output the final segmentation result.

[0008] Furthermore, in step 3, the multi-scale supervision signal is generated by an independent UNet structure, which outputs feature maps at four different stages of its decoder, and then fuses them with the features of the corresponding layer of the main model M through a channel attention mechanism.

[0009] Furthermore, the consistency loss in step 4 is defined as: , in, , This indicates the initial output of the main model M. This indicates a noisy mask. Indicates time-related scaling. This represents the time-related jump weights.

[0010] Furthermore, the segmentation supervision loss in step 4 is defined as: , in, This represents the true segmentation mask, used to directly constrain the accuracy of the model output.

[0011] Furthermore, the overall loss function in step 4 is: , in, This represents the balancing hyperparameter.

[0012] Furthermore, in step 4, the parameter update of the target model TM uses the momentum coefficient. EMA strategy that changes with the training process: , in, express, express.

[0013] Furthermore, in step 5, the inference process does not require iterative denoising; it only performs a single forward propagation, which takes an average of 1.9 seconds.

[0014] Furthermore, in step 2, the continuous time step Calculated using the following formula: , in, , , , express, This indicates that it is used to control noise scheduling and training stability.

[0015] The beneficial effects of this invention are: 1. This invention introduces the consistency model into the medical image segmentation task for the first time, and constructs a segmentation framework (CTS) based on consistency learning. Through the collaborative training mechanism of the main model and the target model, the dependence on multi-step sampling is significantly reduced while ensuring segmentation accuracy, providing a new theoretical path for efficient medical image segmentation. 2. This invention designs a multi-scale supervised signal generation module, which uses an independent UNet structure to extract multi-level contextual features from the original medical image and integrates them into the decoding stage of the consistency model through a channel attention mechanism, effectively enhancing the model's ability to perceive small targets and complex boundaries and improving the detail preservation ability of segmentation. 3. This invention proposes a joint loss function that combines the consistency training loss (LCT) with the direct segmentation supervision loss (LS). During training, it simultaneously optimizes the consistency of output between time steps and the alignment between the predicted mask and the real label, thereby achieving high-precision segmentation even under single-step inference conditions. 4. This invention achieves a highly efficient single-step inference mechanism, which only requires one forward propagation to generate the final segmentation result during the testing phase. The average inference time is only 1.9 seconds, which significantly shortens the prediction time compared with the traditional diffusion segmentation model and significantly improves the feasibility of clinical deployment. 5. The present invention has been systematically validated on multiple publicly available medical image datasets. Experimental results show that the proposed CTS model outperforms existing diffusion models and mainstream segmentation methods in key indicators such as Dice coefficient and IoU. It has both high segmentation accuracy and high inference efficiency, and has good generalization ability and practical value. Attached Figure Description

[0016] Figure 1 This is a schematic diagram of the CTS prediction process.

[0017] Figure 2 This is a schematic diagram of an attention-based signal aggregation method; Figure 3 This is the overall flowchart of the CTS model.

[0018] Example Example 1 Step 1: Construct a medical image dataset, perform standardization preprocessing on the original images and their corresponding segmentation masks, and divide them into training and test sets; This embodiment uses three public datasets for validation: BraTS-2021 (brain tumor MRI), SEHPI (liver tumor ultrasound), and thyroid nodule ultrasound dataset. For the BraTS-2021 dataset, it contains 1251 3D MRI images with an original resolution of 240×240×155. In the preprocessing stage, all 3D images were sliced ​​into 2D images and uniformly resized to 256×256 pixels. At the same time, anisotropic diffusion filtering

[13] was applied to the medical images to remove Poisson noise and preserve edge information. The pixel intensity of all images was normalized to the [0, 1] interval. The dataset was divided into training and testing sets according to standard protocols, with tumor regions manually annotated by experts (such as tumor core, enhanced tumor, edema, etc.) used as segmentation masks.

[0019] Step 2: Select a real segmentation mask from the training set as the target sample, and construct a noisy input in the continuous time domain using Gaussian noise; during training, randomly select a sample pair from the training set. ),in These are preprocessed medical images. It is the corresponding true segmentation mask (i.e., the target sample). Based on the continuous-time setting of the consistency model, =0.002, T=80. (Using the formula...)

[0020] The discrete index is mapped to a continuous time step. In each iteration, a time step is randomly selected. and to the real mask Add Gaussian noise z to construct a noisy input:

[0021] Step 3: Design a multi-scale supervised signal generation module to extract multi-scale contextual features from the original medical image and adaptively fuse them into the decoding path of the consistency master model through a channel attention mechanism to generate preliminary segmentation predictions; this embodiment uses two independent 4× UNet architectures. One of the UNets serves as the multi-scale supervised signal generation module h( This module uses raw medical images. As input, feature maps U are progressively upsampled and generated at four different scales in the four different decoding stages of its encoder-decoder structure.

[0022] (i.e., multi-scale supervision signals). These feature maps contain rich information ranging from global context to local details. Subsequently, these supervision signals are fused with the corresponding decoding layer features of the master consistency model M through a channel attention mechanism. This mechanism automatically learns the importance weight of each channel, thereby more effectively utilizing the key information in the supervision signals to guide the master model in generating more accurate initial segmentation predictions. .

[0023] Step 4: Construct a consistent training framework, including a main model M and a target model TM, where the parameters of TM are updated by the exponential moving average (EMA) of M; calculate the main model at time step... The output and the target model at adjacent time steps The consistency loss between the outputs is combined with the direct segmentation supervision loss to form a joint optimization objective; The training framework consists of two core components: the main model M (another 4× UNet) and the target model TM (with the same structure as M). First, using... Multi-scale supervision signal U and time step As input, the initial output is obtained through the main model M. And calculate its final output:

[0024] in =- and =1. Next, use adjacent time steps. Construct a new noisy input The output is calculated through the target model TM (whose supervision signal is generated by an independent module). Parameters of the target model TM Through the parameters of the main model M Update using the Exponential Moving Average (EMA):

[0025] momentum coefficient The loss function is dynamically adjusted with the number of training iterations k. Ultimately, the joint loss function consists of two parts: Consistency training loss:

[0026] Segmented monitoring loss:

[0027] The overall optimization goal is hyperparameters =1. The model is suitable for end-to-end training using the AdamW optimizer with a batch size of 8.

[0028] Step 5: During the testing phase, only one forward propagation is required: input the pure noise, multi-scale supervision signal and initial time step into the trained master model, and directly output the final segmentation result.

[0029] During the inference phase, the CTS model exhibits extremely high efficiency. First, it samples pure noise once from a standard Gaussian distribution. Then, it processes the original medical image to be segmented. Input the pre-trained supervision signal generation module h( ), thus obtaining the multi-scale supervision signal U Finally, pure noise Supervisory signal U The initial time step t = ε is input along with the pre-trained main model M. After one forward propagation, the model directly outputs the final segmentation result. No iterative denoising process is required. Experiments show that the average time for this single-step inference process is only 1.9 seconds, and it achieves better Dice and IoU scores than existing methods (such as MedSegDiff, TransUNet, etc.) on BraTS-2021, SEHPI, and thyroid nodule datasets, successfully achieving a balance between high accuracy and high efficiency.

[0030] The above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention in any way. Although the present invention has been disclosed above with reference to preferred embodiments, it is not intended to limit the present invention. Any person skilled in the art can make some modifications or alterations to the above-disclosed technical content to create equivalent embodiments without departing from the scope of the present invention. Any simple modifications, equivalent substitutions, and improvements made to the above embodiments without departing from the scope of the present invention, based on the technical essence of the present invention and within the spirit and principles of the present invention, shall still fall within the protection scope of the present invention.

Claims

1. An efficient medical image segmentation method based on consistency learning, characterized in that, The specific steps include: Step 1: Construct a medical image dataset, standardize the original images and their corresponding segmentation masks, and divide them into training and test sets; Step 2: Select real segmentation masks from the training set as target samples, and construct noisy inputs in the continuous time domain using Gaussian noise; Step 3: Design a multi-scale supervised signal generation module to extract multi-scale contextual features from the original medical image and adaptively fuse them into the decoding path of the consistency master model through a channel attention mechanism to generate preliminary segmentation predictions. Step 4: Construct a consistent training framework, including a main model M and a target model TM, where the parameters of TM are updated by the exponential moving average of M; calculate the parameters of the main model M at time step... The output and the target model TM at adjacent time steps The consistency loss between the outputs is combined with the direct segmentation supervision loss to form a joint optimization objective; Step 5: During the testing phase, only one forward propagation is required: input the pure noise, multi-scale supervision signal and initial time step into the main model M to train, and directly output the final segmentation result.

2. The efficient medical image segmentation method based on consistency learning according to claim 1, characterized in that, In step 3, the multi-scale supervision signal is generated by an independent UNet structure, which outputs feature maps at four different stages of its decoder and fuses them with the features of the corresponding layer of the main model M through a channel attention mechanism.

3. The efficient medical image segmentation method based on consistency learning according to claim 1, characterized in that, In step 4, the consistency loss is defined as: ， in, , This indicates the initial output of the main model M. This indicates a noisy mask. Indicates time-related scaling. This represents the time-related jump weights.

4. The efficient medical image segmentation method based on consistency learning according to claim 1, characterized in that, In step 4, the segmentation supervision loss is defined as: ， in, This represents the true segmentation mask, used to directly constrain the accuracy of the model output.

5. A highly efficient medical image segmentation method based on consistency learning according to claim 1, 3, or 4, characterized in that, The overall loss function in step 4 is: ， in, This represents the balancing hyperparameter.

6. The efficient medical image segmentation method based on consistency learning according to claim 1, characterized in that, In step 4, the parameter update of the target model TM uses the momentum coefficient. EMA strategy that changes with the training process: ， in, express, express.

7. The efficient medical image segmentation method based on consistency learning according to claim 1, characterized in that, In step 5, the inference process does not require iterative denoising and only performs a single forward propagation, which takes an average of 1.9 seconds.

8. The efficient medical image segmentation method based on consistency learning according to claim 1, characterized in that, Continuous time steps in step 2 Calculated using the following formula: ， in, , , , express, This indicates that it is used to control noise scheduling and training stability.

Images