A medical image segmentation method based on a conditional diffusion model

Abstract

The application discloses a medical image segmentation method based on a conditional diffusion model, adopts a constructed segmentation network model to segment a medical image, obtains the medical image, inputs the medical image after pretreatment into a Transform encoder to obtain first features, inputs the first features into a conditional network decoder to obtain second features, obtains a noise-added mask image, splices the noise-added mask image and the second features, inputs the noise-added mask image and the second features into a diffusion model encoder to obtain third features, then inputs the third features and the first features into a diffusion model middle layer to obtain fourth features, and finally inputs the fourth features into a diffusion model decoder to obtain a final segmentation result. The application fully combines image and mask information for learning, increases context information in the medical image, better locates a region needing segmentation in the medical image, and thus improves the performance of a deep neural network and the accuracy of a segmentation model.

Description

TECHNICAL FIELD

[0001] The present application relates to the field of artificial intelligence computer vision images, and more particularly to a medical image segmentation method based on a conditional diffusion model. BACKGROUND

[0002] In recent years, with the gradual development of artificial intelligence technology, image segmentation technology has made rapid progress. The purpose of medical image segmentation is to segment the part of the medical image that has certain special meaning and extract relevant features to provide reliable basis for clinical diagnosis and pathological research, and assist doctors in making more accurate diagnosis. It is one of the most challenging tasks in medical image analysis.

[0003] Currently, medical image segmentation faces many problems, including image quality problems: medical images are prone to low-quality images (blur, noise, artifacts, etc.) due to quality problems of acquisition equipment, and it is difficult to obtain annotations: medical image data usually requires experts and doctors with relevant experience, so it is difficult to obtain high-quality annotations. And the problem of large image difference: there are many types of medical images, such as MRI, CT, etc., and the differences between individual cases are also large, causing large image differences, which will also cause difficulties for image analysis.

[0004] Although diffusion models have been proven to generate high-quality images while providing desirable properties such as distribution coverage, smooth training objectives, and ease of extension, there has been relatively little research on diffusion model architectures, and diffusion models still have the potential for further improvement. However, the original diffusion model has three main drawbacks: slow sampling speed, poor maximum likelihood difference, and weak data generalization ability. The generation results of diffusion models have diversity, which makes it difficult to apply directly to medical image segmentation.

[0005] Therefore, it is of great practical significance to study medical image segmentation technology, develop efficient medical image segmentation algorithms, improve the technical level of image segmentation methods, and obtain high-quality segmentation results, which can promote the development of computer vision and medical fields. SUMMARY

[0006] The purpose of the present application is to provide a medical image segmentation method based on a conditional diffusion model to overcome the problems in the technical background.

[0007] To achieve the above-mentioned purpose, the technical scheme of the present application is as follows:

[0008] A medical image segmentation method based on a conditional diffusion model, which adopts a constructed segmentation network model to segment a medical image, the segmentation network model comprising a Transformer encoder, a conditional network decoder, a diffusion model encoder, a diffusion model intermediate layer and a diffusion model decoder, the medical image segmentation method based on the conditional diffusion model comprising:

[0009] Obtaining a medical image, pre-processing and inputting into the Transformer encoder to obtain a first feature;

[0010] Inputting the first feature into the conditional network decoder to obtain a second feature;

[0011] Obtaining a noise-added mask image, performing feature splicing on the noise-added mask image and the second feature, inputting into the diffusion model encoder to obtain a third feature;

[0012] Inputting the third feature and the first feature into the diffusion model intermediate layer to obtain a fourth feature;

[0013] Inputting the fourth feature into the diffusion model decoder to obtain a final segmentation result.

[0014] Further, the pre-processing comprises:

[0015] Converting the medical image into a format, then positioning a region of interest, adjusting the window width and window level to brighten the region of interest, and finally cropping the medical image to a predetermined size.

[0016] Further, the conditional network decoder is composed of three blocks, the first two blocks each containing bilinear interpolation upsampling and double convolution, and the last block including a convolution block and bilinear interpolation upsampling.

[0017] Further, when training the segmentation network model, the noise-added mask image is obtained using the following formula:

[0018]

[0019] Where α t is a hyperparameter, x t represents a noise-added mask image at time t, z t represents noise at time t, and x0 represents a real mask image corresponding to the training image.

[0020] After training the segmentation network model, the noise-added mask image directly uses a pure Gaussian noise image when segmenting the medical image using the trained segmentation network model.

[0021] Further, the diffusion model encoder comprises four encoding layers, each encoding layer having two residual blocks and a linear attention layer, the first three encoding layers further comprising down-sampling, and the last encoding layer further comprising a 1x1 convolution layer.

[0022] The diffusion model decoder comprises four decoding layers, each decoding layer having two residual blocks and a linear attention layer, the first three decoding layers further comprising up-sampling, and the last decoding layer further comprising a 1x1 convolution layer.

[0023] Further, the diffusion model intermediate layer comprises a first residual block, three stacked double cross-attention modules, a self-attention module, and a second residual block.

[0024] The third feature is first input to the first residual block, and then input to the three double cross-attention modules respectively; the first feature is input to the first double cross-attention module; the output of the previous double cross-attention module is input to the next double cross-attention module.

[0025] The output of the third double cross-attention module is further input to the self-attention module and the second residual block, and finally the fourth feature is obtained.

[0026] The double cross-attention module comprises two cross-attention modules,

[0027] In the first cross-attention module of the first double cross-attention module, the first feature is mapped into k and v vectors, and the output of the first residual block is mapped into a q vector; in the second cross-attention module of the first double cross-attention module, the output of the first cross-attention module is mapped into a q vector, and the output of the first residual block is mapped into k and v vectors.

[0028] In the first cross-attention module of the second and third double cross-attention modules, the output of the previous double cross-attention module is mapped into k and v vectors, and the output of the first residual block is mapped into a q vector; in the second cross-attention module of the second and third double cross-attention modules, the output of the first cross-attention module is mapped into a q vector, and the output of the first residual block is mapped into k and v vectors.

[0029] Further, the residual block comprises two blocks, each block comprising a 3x3 convolution, group normalization, and a SiLU activation function, and the first block further comprises time step encoding information before the SiLU for scaling and shifting the input feature.

[0030] The medical image segmentation method based on the conditional diffusion model provided in the application uses the diffusion model of UNet with residual blocks based on the encoding layer and the decoding layer, which can make the network converge better; the conditional neural network decoder structure adopts bilinear interpolation upsampling and double convolution, which can learn an adaptive mapping to recover the features of the medical image with more detailed information. In addition, the model can fully combine image and mask information for learning, increase the context information in the medical image, better locate the region that needs to be segmented in the medical image, thereby improving the performance of the deep neural network and the accuracy of the segmentation model, which also effectively alleviates the problem of small medical image dataset. Finally, in order to obtain more accurate segmentation results, the application combines the features from the Transformer encoder and the diffusion model encoder for feature interaction, so that the model can more accurately segment, which also effectively solves the problem of diversity of the diffusion model generated results. BRIEF DESCRIPTION OF DRAWINGS

[0031] Figure 1 The flow chart of the medical image segmentation method based on the conditional diffusion model of the application.

[0032] Figure 2 The overall framework diagram of the segmentation network model of the application.

[0033] Figure 3 The conditional network decoder schematic diagram of the embodiment of the application.

[0034] Figure 4 The intermediate layer schematic diagram of the diffusion model of the embodiment of the application.

[0035] Figure 5 The DCA module schematic diagram of the embodiment of the application.

[0036] Figure 6 The residual block schematic diagram of the embodiment of the application. DETAILED DESCRIPTION

[0037] In order to make the purpose, technical scheme and advantages of the application clearer, the application will be further described in detail below in combination with the drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the application and not to limit the application.

[0038] The application provides a medical image segmentation method based on a conditional diffusion model, which first constructs a segmentation network model such as Figure 2As shown, the method includes a Transformer encoder, a conditional network decoder, a diffusion model encoder, a diffusion model intermediate layer, and a diffusion model decoder. The model is trained by training image pairs and their corresponding noisy mask images. After the segmentation network model is trained, the medical image to be segmented is input into the trained model to obtain the segmentation result of the medical image.

[0039] In one embodiment, as shown in Figure 1 A medical image segmentation method based on a conditional diffusion model is proposed, as shown in the following formula (1), which includes:

[0040] Step S1, obtain a medical image, and input the preprocessed medical image into a Transformer encoder to obtain first features.

[0041] In this embodiment, whether the segmentation network model is trained or the trained segmentation network model is used for medical image segmentation, the medical image input into the model needs to be preprocessed.

[0042] In the preprocessing process, the medical image is converted from a CT image format to a JPG format, then the region of interest (ROI) is located, the window width and window center of the ROI are adjusted, and finally the medical image is cropped to 192x256.

[0043] The window technology in the field of medical images includes window width and window center, which are used to select the CT value range of interest. Because various tissue structures or lesions have different CT values, when displaying a certain tissue structure detail, the window width and window center suitable for observing the tissue or lesion should be selected to obtain the best display. The window width and window center correspond to the contrast and brightness of the image, respectively. The required processing of the region for adjusting the window width and window center is not described here.

[0044] After preprocessing, the features become 192x256, and the lesion area (i.e., the region of interest) is highlighted, so that the network can well locate the lesion.

[0045] The Transformer encoder in this embodiment includes 8 Transformer blocks, each of which uses a 16x16 patch Vision Transformer, which can fully model the global information of the medical image.

[0046] Before the medical image is fed into the Transformer encoder, it is first divided into 16×16 patches using a Patch Embedding layer, dividing the medical image into n 16×16 feature blocks. This is beneficial for network training as it reduces the number of parameters. Then, a linear embedding layer converts the features into tokens required by the Transformer, which are finally input into an 8-layer Transformer encoder to model and learn the global information of the features.

[0047] Step S2: Input the first feature into the conditional network decoder to obtain the second feature.

[0048] The conditional network decoder in this embodiment is as follows: Figure 3 As shown, it consists of three blocks. The first two blocks each contain a bilinear interpolation upsampling and a double convolution module. Upsampling increases the image size, and double convolution obtains higher resolution features, which can better model the feature information. The last block contains a convolution block and a bilinear interpolation upsampling. This convolution block consists of a 3×3 convolution, a BN layer, a ReLU layer, and a 3×3 convolution.

[0049] The first two blocks upsample the feature information by a factor of two to recover high-resolution features. Finally, there is a convolutional block and a four-fold bilinear interpolation upsampling. This convolutional block processes the feature before output, resulting in more detailed features.

[0050] Step S3: Obtain the noisy mask image, concatenate the noisy mask image with the second feature, and input it into the diffusion model encoder to obtain the third feature.

[0051] The formula for obtaining the noisy mask image during the training of the segmentation network model in this application is as follows:

[0052]

[0053] Where α t It is a hyperparameter, x t Z represents the noisy mask image at time t. t Then, it represents the noise at time t. This formula can be used to obtain the noise-added mask image at any time. x0 represents the real mask image corresponding to the training image.

[0054] Add noise masking image x t The second feature is concatenated along the channel dimension, and then fed into the diffusion model encoder along with the time step t. The time step t is mapped to the same size as the feature through a multilayer perceptron, and then the input feature is scaled and shifted in each encoder layer.

[0055] After the segmentation network model is trained, the trained segmentation network model is used to segment the medical image, and the noise mask image directly uses a pure Gaussian noise image.

[0056] The diffusion model encoder of the embodiment includes four encoding layers, each of which has two residual blocks and a linear attention layer, and the first three encoding layers further include down-sampling, and the last encoding layer further includes a 1x1 convolution layer.

[0057] As shown in Figure 6 each residual block (ResnetBlock) includes two blocks, each block is composed of a 3x3 convolution, a group normalization and a SiLU activation function, but the first block of each ResnetBlock further adds time step encoding information to scale and shift the input features before SiLU.

[0058] The linear attention layer models the context information of the input features, enhancing the detailed information of the features.

[0059] The structure of the residual block of the embodiment is applicable to the diffusion model encoder, the intermediate layer and the decoder, and the structures of all residual blocks are the same, which will not be described below.

[0060] Step S4, input the third feature and the first feature into the diffusion model intermediate layer to obtain a fourth feature.

[0061] The diffusion model intermediate layer of the embodiment includes a first residual block, three stacked double cross attention modules (DCA), a self-attention module and a second residual block, as shown in Figure 4 Each DCA is composed of two cross attentions, and each DCA module interacts with the input features to enable the noise image to learn the key information in the training image. Figure 5

[0062] Specifically, the third feature is first passed through the first residual block, and then input into three double cross attention modules; the first feature is input into the first double cross attention module; the output of the previous double cross attention module is input into the next double cross attention module. The output of the third double cross attention module is further passed through the self-attention module and the second residual block, and finally the fourth feature is obtained.

[0063] The DCA module includes two cross attentions, as shown in Figure 5

[0064] ​​In the first cross attention of the first double cross attention module, the first feature is mapped as a k and v vector, and the output of the first residual block is mapped as a q vector; in the second cross attention of the first double cross attention module, the output of the first cross attention is mapped as a q vector, and the output of the first residual block is mapped as a k and v vector.

[0065] In the first cross attention of the second and third double cross attention modules, the output of the previous double cross attention module is mapped as a k and v vector, and the output of the first residual block is mapped as a q vector; in the second cross attention of the second and third double cross attention modules, the output of the first cross attention is mapped as a q vector, and the output of the first residual block is mapped as a k and v vector.

[0066] In the cross attention, the q feature and the k feature are first multiplied, then activated by a sigmoid activation function, and finally multiplied by the v feature mapping, and finally a time step t is added for shifting and scaling.

[0067] After the DCA module, the noise image and the training image have a good semantic interaction, and the features are well enhanced, and then the self-attention layer and the residual block are used to model and enhance the features themselves.

[0068] Step S5, input the fourth feature into the diffusion model decoder to obtain the final segmentation result.

[0069] The diffusion model decoder of the embodiment is similar to the diffusion model encoder, except that the encoder is to upsample and downsize the features, and the decoder is to downsample and restore the feature size.

[0070] The diffusion model encoder and the diffusion model decoder of the embodiment include four decoding layers, each decoding layer has two residual blocks and a linear attention layer, the first three decoding layers further include an upsampling, and the last decoding layer further includes a 1x1 convolution layer. The input feature of each decoding layer is first spliced with the output feature of the corresponding encoding layer and then input into the decoding layer.

[0071] The fourth feature is input into the diffusion model decoder, and the decoder and the encoder have similar structures, except that the down-sampling module is replaced by the up-sampling module, and finally the segmentation mask is output.

[0072] In one specific embodiment, the loss function of the segmentation network model of the application is as follows:

[0073]

[0074] Where w is a weight matrix, and the loss is composed of two parts, wherein the cross entropy loss function is calculated as follows:

[0075]

[0076] where y i represents the value in the labeled image, f i (x i ) represents the value in the predicted image, and n represents the number of samples. Then is expressed as sum(w x L BCE ) / sum(w).

[0077] And another IOU loss function calculation formula is as follows:

[0078] ① inter = |P∩G| x w

[0079] ② union = |P∪G| x w

[0080] ③

[0081] Wherein, P represents the predicted image mask, and G represents the real image mask.

[0082] During training, the output of the diffusion model decoder is passed through a sigmoid activation function and the real mask image is put into the loss function to calculate the loss until the loss of the segmentation model no longer converges, completing the training of the segmentation network model.

[0083] After training the segmentation network model, load the trained model parameters, input the medical image into the trained deep neural network for image segmentation. When image segmentation, from a time step t, x t is a pure Gaussian noise image, the input of the Transformer encoder is the medical image, and the result of the medical image after passing through the Transformer encoder and the conditional network decoder is y, x t and y are spliced and used as the input of the diffusion model to predict the image x t-1 at the previous moment, then x t-1 is spliced with y and input into the diffusion model to predict the image x t-2 , step by step to restore the result at the last moment 0 as the final segmentation result.

[0084] The above-described embodiments only express several embodiments of the present application, and the description is more specific and detailed, but it cannot be understood as limiting the scope of the patent. It should be noted that for ordinary skilled in the art, without departing from the concept of the present application, a number of modifications and improvements can be made, which are within the scope of the present application. Therefore, the scope of protection of the patent of the present application should be subject to the appended claims.

Claims

1. A medical image segmentation method based on a conditional diffusion model, using a constructed segmentation network model to segment a medical image, characterized in that, The segmentation network model comprises a Transformer encoder, a conditional network decoder, a diffusion model encoder, a diffusion model intermediate layer and a diffusion model decoder, and the medical image segmentation method based on the conditional diffusion model comprises the following steps: Obtain a medical image, input the preprocessed medical image into the Transformer encoder to obtain first features; Input the first features into the conditional network decoder to obtain second features; Obtain a noise-added mask image, splice the noise-added mask image and the second features to obtain third features, and input the third features into the diffusion model encoder to obtain fourth features; Input the third features and the first features into the diffusion model intermediate layer to obtain the fourth features; the diffusion model intermediate layer comprises a first residual block, three stacked double cross-attention modules, a self-attention module and a second residual block; the third features pass through the first residual block first, and then are input into the three double cross-attention modules respectively; the first features are input into the first double cross-attention module; the output of the previous double cross-attention module is input into the next double cross-attention module; the output of the third double cross-attention module passes through the self-attention module and the second residual block to obtain the fourth features finally; Input the fourth features into the diffusion model decoder to obtain a final segmentation result. 2.The method of claim 1, wherein, The preprocessing comprises the following steps: Convert the format of the medical image, locate a region of interest, adjust the window width and window level of the region of interest, and finally crop the medical image to a predetermined size. 3.The method of claim 1, wherein, The conditional network decoder is composed of three blocks, the first two blocks each comprise bilinear interpolation upsampling and double convolution, and the last block comprises a convolution block and bilinear interpolation upsampling. 4.The method of claim 1, wherein, When training the segmentation network model, the noise-added mask image is obtained by using the following formula: ； wherein is a hyperparameter, denotes a noisy mask image at time t, denotes noise at time t, denotes a real mask image corresponding to the training image; After the segmentation network model is trained, when the trained segmentation network model is used to segment the medical image, the noise-added mask image directly uses a pure Gaussian noise image. 5.The method of claim 1, wherein, The diffusion model encoder comprises four encoding layers, each encoding layer has two residual blocks and a linear attention layer, the first three encoding layers further comprise down-sampling, and the last encoding layer further comprises a 1x1 convolution layer; The diffusion model decoder comprises four decoding layers, each decoding layer has two residual blocks and a linear attention layer, the first three decoding layers further comprise up-sampling, and the last decoding layer further comprises a 1x1 convolution layer. 6.The method of claim 1, wherein, The double cross-attention module comprises two cross-attentions, In the first cross-attention of the first double cross-attention module, the first features are mapped into k and v vectors, and the output of the first residual block is mapped into a q vector; in the second cross-attention of the first double cross-attention module, the output of the first cross-attention is mapped into a q vector, and the output of the first residual block is mapped into k and v vectors; In the first cross-attention of the first double cross-attention module, the first features are mapped into k and v vectors, and the output of the first residual block is mapped into a q vector; in the second cross-attention of the first double cross-attention module, the output of the first cross-attention is mapped into a q vector, and the output of the first residual block is mapped into k and v vectors; In the first cross-attention of the second and third double cross-attention modules, the output of the previous double cross-attention module is mapped into k and v vectors, and the output of the first residual block is mapped into a q vector; in the second cross-attention of the second and third double cross-attention modules, the output of the first cross-attention is mapped into a q vector, and the output of the first residual block is mapped into k and v vectors.

7. The medical image segmentation method based on a conditional diffusion model according to claim 5 or 6, characterized in that, The residual block includes two blocks, each block is composed of a 3*3 convolution, a group normalization and a SiLU activation function, and the first block further adds time step encoding information to the input features before the SiLU for scaling and shifting.

Images