A medical image segmentation method based on a pre-trained diffusion model

Abstract

The application discloses a medical image segmentation method based on a pre-training diffusion model, which comprises the following steps: 1, obtaining a medical image segmentation dataset and preprocessing; 2, establishing a forward noise adding, reverse noise removing process and a pre-training prediction network; and 3, training a segmentation network and a prediction. The application directly predicts an original label and prior knowledge of a large model through a diffusion network, and improves the segmentation performance of a target region in a medical image.

Description

Technical Field

[0001] This invention relates to the field of image signal processing technology, and specifically to a medical image segmentation method based on a pre-trained diffusion model. Background Technology

[0002] With the rapid development of digital healthcare technology, medical image segmentation has become a crucial step in clinical auxiliary diagnosis, surgical planning, and quantitative analysis of lesions. Early medical image segmentation primarily relied on convolutional neural networks (CNNs), such as U-Net and its variants, which extracted image features through an encoder-decoder structure. In recent years, Transformer-based models have demonstrated excellent performance in handling complex anatomical structures by utilizing self-attention mechanisms to capture global contextual information. However, due to the presence of significant imaging noise, artifacts, and blurred target region boundaries in medical images (such as CT and MRI), traditional segmentation networks often suffer from decreased segmentation accuracy, resulting in edge loss or oversegmentation, when dealing with subtle edges and complex background interference. Although pre-trained large-scale visual models (such as SAM) possess powerful general feature representation capabilities, directly fine-tuning them to adapt to medical imaging tasks is not only computationally expensive but may also lead to the model losing its original general spatial understanding capabilities.

[0003] Diffusion models, as a novel generative model, have achieved significant results in image processing through iterative denoising processes. Although existing research has attempted to apply diffusion models to medical image segmentation, they often suffer from problems such as numerous iterations, slow convergence speed, and limited feature representation when processing high-dimensional medical data. Particularly in the inverse denoising process, the lack of effective modeling of medical-specific manifold features leads to insufficient sensitivity in responding to lesion regions. Therefore, how to combine the prior knowledge of large-scale pre-trained models to improve the segmentation accuracy of medical images in noisy environments while reducing computational costs has become a pressing technical problem in the field. Summary of the Invention

[0004] To overcome the problems existing in medical image segmentation of current image segmentation techniques, this invention provides a medical image segmentation method based on a pre-trained diffusion model, which aims to improve segmentation accuracy. This method is less susceptible to imaging noise interference and the easy loss of lesion edge information, thereby improving the accuracy of image segmentation and providing support for further image analysis and processing.

[0005] To solve the above problems, the present invention adopts the following technical solution: The medical image segmentation method based on a pre-trained diffusion model of the present invention is characterized by the following steps: Step 1: Obtain public medical image datasets ,in, For the first A label image, For the first There are N public medical images, where N is the total number of public medical images; Step 2, The input is processed in the forward noise-adding module, thereby... Within each discrete time step, standard Gaussian noise is injected incrementally using equation (1). , obtained the A sequence of noisy images ,in, express The Noisy images at each time step: (1) In equation (1), To be standard Gaussian noise that follows a normal distribution, It is the first The cumulative product coefficients at each time step, and ,in, Indicates the first Noise figure at each time step; Step 3: Construct the reverse noise reduction module and to and Processing is performed to obtain The Estimates of the label image at each time step ; Step 4: Construct the multi-task combined loss function of the segmentation network consisting of a forward noise addition module and a backward noise reduction module. ; Step 5: Iteratively train the segmentation network using the AdamW optimizer and calculate the multi-task combination loss function. Adjust network parameters until the multi-task combined loss function is reached. The process continues until convergence, thus obtaining a well-trained segmentation model. Step 6: In the inference phase, random samples of pure noise tensors are taken from a standard Gaussian distribution and input along with the medical images into the trained segmentation model, thereby... After denoising at each time step, the predicted label is obtained and used as the segmentation result.

[0006] The medical image segmentation method based on a pre-trained diffusion model described in this invention is also characterized in that step 3 includes: Step 3.1: Initialize t=T; define the sequence number of the denoising time step as... Where S is the sampling step size; Step 3.2 right and Process and output. The Estimates of the label image at each time step ; Step 3.3: Calculate the first step using equation (2). Noisy images at each time step ; (2) In equation (2), for The mean function, For the first The variance of the noise figure at each time step This represents noise that follows a normal distribution. Step 3.4: After assigning T-1 to T, return to step 3.2 and execute sequentially until... = Until then, thus obtaining The Estimates of the label image at each time step .

[0007] Furthermore, in step 3.2 include: Layer label encoder Layer image encoder, Layer decoder; Step 3.2.1: Each label encoder layer includes: a self-attention sublayer and a hybrid expert sublayer; when =1, No. Self-attention sublayer pairs in layer label encoder and After performing layer normalization and self-attention processing, the first layer is obtained. The first time step Intermediate features of the layer ;Will The input is fed into the hybrid expert sublayer and passes through the parameter-frozen shared expert module, gated router, and... The weighted mapping process of the routing expert outputs the first... The first time step Layer coding features ; when =2,3,…, At that time, the first The first time step Layer coding features Enter the first The processing is carried out in the layer label encoder, thereby gradually moving from the first layer to the second layer. Layer label encoder output The first time step Layer coding features ; Step 3.2.2 Layer image encoder sequentially Perform convolution processing to obtain Layer coding features ,in, Indicates the first The coding characteristics of the layer; Step 3.2.3, will and Enter the first L Upsampling and convolution are performed in the layer decoder to obtain the first... Decoding features of the Lth layer at each time step ; when = When L-1, L-2, ..., 1, the first The first time step Layer decoding features , and Enter the first The process is carried out in the layer decoder, thus gradually outputting the next layer from the layer 1 decoder. The first time step Layer coding features And then Perform a convolution operation with a 1×1 kernel to generate... .

[0008] Furthermore, the hybrid expert sublayer in step 3.2.1 is obtained according to the following steps. : Step a, to In the shared expert module where input parameters are frozen, and using a feedforward neural network layer... After transformation, the output is the first... The first time step Layer shared features ; Step b, will The input is processed in the gating router to obtain the first... The first time step Layer weight vector and to In Sort the weights in descending order. and from Before the election The target expert corresponding to each weight; among them. For the first The first time step Layer weight vector The Each weight; Step c, will Enter them separately Process the target experts and output the first one. The first time step layer Expert characteristics ;in, Indicates the first The first time step Layer One expert characteristic; Step d, will Expert characteristics Separately with the previous Weights Corresponding products are multiplied, and the resulting k products are then multiplied by the shared feature. After accumulation, output the first... The first time step Layer coding features .

[0009] Furthermore, the steps include: Step 4.1: Construct the Dice loss using equation (3). : (3) In equation (3), Indicates the smoothing term; Step 4.2: Construct the binary cross-entropy loss using equation (4). : (4) Step 4.3: Construct the mean squared error loss using equation (5). : (5) In equation (5), express and Mathematical expectation under joint distribution; Step 4.5: Construct the multi-task combination loss function using equation (6). ; (6).

[0010] The present invention provides an electronic device, including a memory and a processor, characterized in that the memory is used to store a program that supports the processor in executing the medical image segmentation method, and the processor is configured to execute the program stored in the memory.

[0011] The present invention provides a computer-readable storage medium on which a computer program is stored, characterized in that the computer program, when run by a processor, executes the steps of the medical image segmentation method.

[0012] Compared with existing segmentation methods, the beneficial effects of this invention are as follows: 1. This invention significantly improves the segmentation accuracy and structural consistency of medical images in noisy environments by directly predicting the original labeled image, enabling the model to perceive the global contour of the target in the early stages of inverse denoising. Compared to indirect methods of predicting noise, direct prediction... It can more effectively maintain the topological coherence of the target region, significantly reducing edge fragmentation and missegmentation points. At the same time, by utilizing the iterative characteristics of the diffusion model, the model has a strong tolerance for artifacts and imaging noise in medical images.

[0013] 2. This invention employs a pre-trained prediction network based on a SAM2 encoder freezing strategy, which, combined with the diffusion model, produces significant synergistic gains. This invention fully leverages the inherent strong generalization space prior and zero-shot learning capabilities of the frozen SAM2 encoder, providing the model with a mature foundation for understanding geometric structures. This allows the diffusion model to no longer solely rely on large-scale medical labeled data to learn complex anatomical shapes from scratch. Furthermore, this stable prior foundation provides high-quality feature representation anchors for the extremely challenging inverse denoising process of the diffusion model, fundamentally avoiding parameter drift or training collapse caused by full fine-tuning when processing noise of different magnitudes during iterations, significantly improving the convergence stability of generative segmentation tasks.

[0014] 3. This invention achieves extremely high parameter utilization and domain adaptation efficiency. It freezes the original feedforward neural network (FFN) layer of the SAM2 encoder as a shared expert in the MoE module, and only specifically fine-tunes the routed experts and lightweight modules. It fully utilizes the general feature extraction capabilities and geometric space priors accumulated by the SAM2 model in large-scale visual pre-training, enabling the model to serve as a stable feature base, greatly reducing computational resource consumption and shortening the adaptation cycle from a general model to a medical-specific model. Through a gated router, medical features are dynamically filtered, achieving deep coupling between the general visual common sense maintained by the shared experts and the medical image-specific tissue texture and lesion distribution dynamically learned by the routed experts, while maintaining extremely low fine-tuning overhead. Attached Figure Description

[0015] Figure 1 This is a flowchart of the medical image segmentation method based on a pre-trained prediction network diffusion model according to the present invention; Figure 2 This is a schematic diagram of the pre-trained prediction network structure implemented in a specific embodiment of the present invention; Figure 3 This is a schematic diagram of the coding block structure implemented in a specific embodiment of the present invention. Detailed Implementation

[0016] In this embodiment, as Figure 1 As shown, a medical image segmentation method based on a pre-trained prediction network diffusion model includes the following steps: Step 1: Obtain a public medical image segmentation dataset to obtain the image dataset. ,in For the first A label image, For the first There are N public medical images, where N is the total number of public medical images; the public medical image segmentation dataset in this embodiment is obtained from Kvasir and ClinicDB. and Performs size normalization, vertical flipping, tensor transformation, and intensity normalization based on predefined mean and standard deviation.

[0017] Step 2, The input is processed in the forward noise-adding module, thereby... Within each discrete time step, standard Gaussian noise is injected incrementally using equation (1). , obtained the A sequence of noisy images ,in, express The Noisy images at each time step: (1) In equation (1), To be standard Gaussian noise that follows a normal distribution, It is the first The cumulative product coefficients at each time step, and ,in, Indicates the first The noise figure at each time step.

[0018] In practice, the forward noise addition process transforms the original labeled image into a purely noisy image through T=1000 steps. The forward noise addition process employs linear noise scheduling, with a noise coefficient... from =0.0001 linearly increases to =0.02.

[0019] Step 3: Construct the reverse noise reduction module ,like Figure 2 As shown, and for and Processing is performed to obtain The Estimates of the label image at each time step .

[0020] Step 3.1: Define the sequence number of the denoising time step as... ,and Where S is the sampling step size, the sampling step size is 20, and the initialization is... .

[0021] Step 3.2 right and Process and output. The Estimates of the label image at each time step .

[0022] Step 3.2.1: Each layer of the label encoder includes: a self-attention sublayer and a hybrid expert sublayer, such as... Figure 3 As shown, the label encoder part uses the SAM2 model encoder pre-trained on large-scale visual data. The frozen part parameters use its large model parameters. The layer normalization and self-attention parameters in the self-attention sub-layer are frozen, and the layer normalization and shared expert parameters in the hybrid expert sub-layer are frozen. Gating and routing experts participate in the training.

[0023] when =1, No. Self-attention sublayer pairs in layer label encoder and After performing layer normalization and self-attention processing, the first layer is obtained. The first time step Intermediate features of the layer ;Will The input is fed into the hybrid expert sublayer and passes through the parameter-frozen shared expert module, gated router, and... The weighted mapping process of the routing expert outputs the first... The first time step Layer coding features .

[0024] when =2,3,…, At that time, the first The first time step Layer coding features Enter the first The processing is carried out in the layer label encoder, thereby gradually moving from the first layer to the second layer. Layer label encoder output The first time step Layer coding features ; Step a, to In the shared expert module where input parameters are frozen, and using a feedforward neural network layer... After transformation, the output is the first... The layer in the first Shared features of each time step ; Step b, will Input the gated router to get the first The first time step Layer weight vector ,right In Sort the weights in descending order. ,from Before being elected The target expert corresponding to each weight; where... For the first The first time step Layer weight vector The Each weight; In the specific implementation, M is set to 3 and k is set to 1. The gated router is implemented by a linear layer and softmax. The routing experts adopt a bottleneck structure design of "dimensionality reduction-nonlinear activation-dimensionality increase", that is, the linear layer is first reduced to 32, then GELU nonlinear activation is performed, and then a linear layer is used to restore the original channel.

[0025] M=3 provides a larger model capacity than M=2, enabling it to capture richer feature patterns; simultaneously, it is easier to optimize than M=4, avoiding the routing collapse problem caused by too many experts on a given dataset. With a moderate total number of experts (3), k=1 acts similarly to hard selection, forcing the model to choose the most suitable single expert for each input. Each expert learns more independent and discriminative features, avoiding information redundancy or averaging that might occur among multiple experts. In current image segmentation tasks, overly fragmented expert combinations (such as M=4, k=2) may require fusing features from different experts in later stages. Since fusion strategies require specialized design, simple fusion strategies may not be as effective as directly using a highly focused expert.

[0026] Step c, will Enter them separately Process the target experts and output the first one. The first time step layer Expert characteristics .in, Indicates the first The first time step Layer A characteristic of an expert.

[0027] Step d, will Expert characteristics respectively with Weights After multiplication, then with shared features After accumulation, output the first... The first time step Layer coding features .

[0028] Step 3.2.2 Layer image encoder sequentially Convolution processing is performed, and the kernel size of the two-dimensional convolutional layer is... With a step size of 1, we obtain Layer coding features ,in, Indicates the first The coding characteristics of the layer; Step 3.2.3, will and Enter the first L The layer decoder performs upsampling and convolution processing up to the 1st... Decoding features of the Lth layer at each time step .

[0029] when = When L-1, L-2, ..., 1, the first The first time step Layer decoding features , and Enter the first The process is carried out in the layer decoder, thus gradually outputting the next layer from the layer 1 decoder. The first time step Layer coding features And then Perform a convolution operation with a 1×1 kernel to generate... .

[0030] Step 3.3: Calculate the first step using equation (2). Noisy images at each time step ; (2) In equation (2), for The mean function, For the first The variance of the noise figure at each time step This represents noise that follows a normal distribution.

[0031] Step 3.4: After assigning T-1 to T, return to step 3.2 and execute sequentially until... = Until then, thus obtaining The Estimates of the label image at each time step .

[0032] Step 4: Construct the multi-task combined loss function of the segmentation network consisting of a forward noise addition module and a backward noise reduction module. ; Step 4.1: Construct the Dice loss using equation (3). : (3) In equation (3), This indicates the smoothing term.

[0033] Step 4.2: Construct the binary cross-entropy loss using equation (4). : (4) Step 4.3: Construct the mean squared error loss using equation (5). : (5) In equation (5), Represents the true value Compared with the predicted value Mathematical expectation under joint distribution.

[0034] Step 4.5: Construct the multi-task combination loss function using equation (6). ; (6) Step 5: Iteratively train the segmentation network using the AdamW optimizer and calculate the multi-task combination loss function. Adjust network parameters until the multi-task combined loss function is reached. The process continues until convergence, thus obtaining a well-trained segmentation model.

[0035] In practice, the AdamW optimizer is used with an initial learning rate of 0.001, a cosine annealing learning rate scheduler is used, the weight decay is 5e-4, and the batch size is 12.

[0036] Step 6: In the inference phase, random samples of pure noise tensors are taken from a standard Gaussian distribution and input along with the medical images into the trained segmentation model, thereby... After denoising at each time step, the predicted label is obtained and used as the segmentation result.

[0037] In this embodiment, an electronic device includes a memory and a processor. The memory stores a program that supports the processor in executing the above-described method, and the processor is configured to execute the program stored in the memory.

[0038] In this embodiment, a computer-readable storage medium stores a computer program, which is executed by a processor to perform the steps of the above method.

[0039] In this embodiment, the present invention compares the proposed method (Ours) with four state-of-the-art methods, including U-Net, PraNet, SAM2_UNet, and SegDiff, on four major public datasets in the field of colon polyp segmentation (Kvasir, ClinicDB, ColonDB, and CVC-300), using mDice (mean dice coefficient) and mIoU (mean intersection-over-union ratio) as the core evaluation metrics. Specifically, U-Net is from the paper "U-Net: Convolutional Networks for Biomedical Image Segmentation", PraNet is from the paper "Pranet: Parallel Reverse Attention Network for Polyp Segmentation", SAM2_UNet is from the paper "SAM2-UNet: Segment Anything 2 Makes Strong Encoder for Natural and Medical Image Segmentation", and SegDiff is from the paper "SegDiff: Image Segmentation with Diffusion Probabilistic Models".

[0040] Please refer to Table 1 for specific results.

[0041] Table 1: Comparison of Quantitative Results of Polyp Dataset

[0042] As shown in Table 1, the proposed method (Ours) achieves state-of-the-art performance across all datasets and evaluation metrics. Specifically: on the relatively well-distributed Kvasir dataset, Ours achieves an mDice of 0.924 and an mIoU of 0.876, representing improvements of 2.55% and 1.98% respectively compared to the suboptimal model SAM2_UNet; on the ClinicDB dataset, Ours slightly outperforms PraNet (0.899) with an mDice of 0.901, and its mIoU (0.848) is roughly on par with PraNet (0.849), maintaining state-of-the-art performance while validating the method's stability; for the ColonDB dataset, which features the most complex polyp morphology and the highest segmentation difficulty, Ours achieves mDice and mIoU of 0.799 and 0.722 respectively, still showing improvements of 0.76% and 0.56% compared to the suboptimal SAM2_UNet. The improvement is significantly better than other comparative models, demonstrating the strong segmentation ability of our method for low-quality, complex-shaped polyps. On the CVC-300 dataset, Ours's mDice (0.892) and mIoU (0.829) are the highest among all models, with improvements of 0.68% and 0.61% respectively compared to SAM2_UNet, further verifying the comprehensiveness of the performance.

[0043] Overall, the proposed method significantly outperforms UNet, PraNet, SAM2_UNet, and SegDiff on four polyp segmentation datasets, demonstrating its effectiveness in feature extraction, edge refinement, and complex target segmentation, and better meeting the clinical needs of medical image segmentation.

Claims

1. A medical image segmentation method based on a pre-trained diffusion model, characterized in that, Follow these steps: Step 1: Obtain public medical image datasets ,in, For the first A label image, For the first There are N public medical images, where N is the total number of public medical images; Step 2, The input is processed in the forward noise-adding module, thereby... Within each discrete time step, standard Gaussian noise is injected incrementally using equation (1). , obtained the A sequence of noisy images ,in, express The Noisy images at each time step: (1) In equation (1), To be standard Gaussian noise that follows a normal distribution, It is the first The cumulative product coefficients at each time step, and ,in, Indicates the first Noise figure at each time step; Step 3: Construct the reverse noise reduction module and to and Processing is performed to obtain The Estimates of the label image at each time step ; Step 4: Construct the multi-task combined loss function of the segmentation network consisting of a forward noise addition module and a backward noise reduction module. ; Step 5: Iteratively train the segmentation network using the AdamW optimizer and calculate the multi-task combination loss function. Adjust network parameters until the multi-task combined loss function is reached. The process continues until convergence, thus obtaining a well-trained segmentation model. Step 6: In the inference phase, random samples of pure noise tensors are taken from a standard Gaussian distribution and input along with the medical images into the trained segmentation model, thereby... After denoising at each time step, the predicted label is obtained and used as the segmentation result.

2. The medical image segmentation method based on a pre-trained diffusion model according to claim 1, characterized in that, Step 3 includes: Step 3.1: Initialize t=T; define the sequence number of the denoising time step as... Where S is the sampling step size; Step 3.2 right and Process and output. The Estimates of the label image at each time step ; Step 3.3: Calculate the first step using equation (2). Noisy images at each time step ; (2) In equation (2), for The mean function, For the first The variance of the noise figure at each time step This represents noise that follows a normal distribution. Step 3.4: After assigning T-1 to T, return to step 3.2 and execute sequentially until... = Until then, thus obtaining The Estimates of the label image at each time step .

3. The medical image segmentation method based on a pre-trained diffusion model according to claim 2, characterized in that, In step 3.2 include: Layer label encoder Layer image encoder, Layer decoder; Step 3.2.1: Each label encoder layer includes: a self-attention sublayer and a hybrid expert sublayer; when =1, No. Self-attention sublayer pairs in layer label encoder and After performing layer normalization and self-attention processing, the first layer is obtained. The first time step Intermediate features of the layer ;Will The input is fed into the hybrid expert sublayer and passes through the parameter-frozen shared expert module, gated router, and... The weighted mapping process of the routing expert outputs the first... The first time step Layer coding features ; when =2,3,…, At that time, the first The first time step Layer coding features Enter the first The processing is carried out in the layer label encoder, thereby gradually moving from the first layer to the second layer. Layer label encoder output The first time step Layer coding features ; Step 3.2.2 Layer image encoder sequentially Perform convolution processing to obtain Layer coding features ,in, Indicates the first The coding characteristics of the layer; Step 3.2.3, will and Enter the first L Upsampling and convolution are performed in the layer decoder to obtain the first... Decoding features of the Lth layer at each time step ; when = When L-1, L-2, ..., 1, the first The first time step Layer decoding features , and Enter the first The process is carried out in the layer decoder, thus gradually outputting the next layer from the layer 1 decoder. The first time step Layer coding features And then Perform a convolution operation with a 1×1 kernel to generate... .

4. The medical image segmentation method based on a pre-trained diffusion model according to claim 3, characterized in that, The hybrid expert sublayer in step 3.2.1 is obtained according to the following steps. : Step a, to In the shared expert module where input parameters are frozen, and using a feedforward neural network layer... After transformation, the output is the first... The first time step Layer shared features ; Step b, will The input is processed in the gating router to obtain the first... The first time step Layer weight vector and to In Sort the weights in descending order. and from Before being selected The target expert corresponding to each weight; among them. For the first The first time step Layer weight vector The Each weight; Step c, will Enter them separately Process the target experts and output the first one. The first time step layer Expert characteristics ;in, Indicates the first The first time step Layer One expert characteristic; Step d, will Expert characteristics Separately with the previous Weights Corresponding products are multiplied, and the resulting k products are then multiplied by the shared feature. After accumulation, output the first... The first time step Layer coding features .

5. The medical image segmentation method based on a pre-trained diffusion model according to claim 4, characterized in that, Step 4 includes: Step 4.1: Construct the Dice loss using equation (3). : (3) In equation (3), Indicates the smoothing term; Step 4.2: Construct the binary cross-entropy loss using equation (4). : (4) Step 4.3: Construct the mean squared error loss using equation (5). : (5) In equation (5), express and Mathematical expectation under joint distribution; Step 4.5: Construct the multi-task combination loss function using equation (6). ; (6)。 6. An electronic device, comprising a memory and a processor, characterized in that, The memory is used to store a program that supports the processor in executing any of the medical image segmentation methods of claims 1-5, and the processor is configured to execute the program stored in the memory.

7. A computer-readable storage medium storing a computer program thereon, characterized in that, The computer program, when run by a processor, performs the steps of any of the medical image segmentation methods described in claims 1-5.

Images