Brain tumor image segmentation method and system based on diffusion model

Abstract

The invention discloses a brain tumor image segmentation method and system based on a diffusion model, and relates to the technical field of brain tumor image segmentation. According to the method, the diffusion model is used for learning the probability distribution of the brain tumor image, effective feature extraction and image segmentation are carried out on the brain tumor image, and accurate brain tumor automatic segmentation is realized. A three-dimensional axial attention mechanism module is introduced into a denoising encoder to construct global remote dependence, so that the capability of capturing context information of the model is remarkably improved, the limitation of convolution is effectively solved, and the interpretability of the model is enhanced. Multi-scale feature fusion is used in a denoising decoder, features of different scales are adaptively fused, rich texture and context information in the multi-scale features are effectively utilized, the feature extraction capacity of the model is enhanced, and therefore the accuracy of brain tumor image segmentation is improved.

Description

Technical Field

[0001] The present invention relates to the technical field of brain tumor image segmentation, and in particular to a brain tumor image segmentation method and system based on a diffusion model. Background Art

[0002] Brain tumors have a high mortality rate in all age groups and rank third in the list of cancer deaths. With the advent of the era of artificial intelligence in the medical field, computer-aided diagnosis and surgical pre-evaluation of brain tumors have become a necessity. Multimodal magnetic resonance imaging (MRI) is often used in clinical practice to provide rich supplementary information for the image segmentation of brain tumors. Therefore, it is of great significance to automatically and accurately segment brain tumors from multimodal MRIs.

[0003] In recent years, deep learning technology has made great progress in brain tumor image segmentation, among which the main technologies include convolutional neural networks (CNNs) and transformers. However, the presence of noise in medical images affects the ability of convolution to extract features. To solve this problem, some existing methods have introduced diffusion models to improve the smoothness of images and reduce noise to improve the quality and accuracy of segmentation results. These models are combined with CNNs to effectively remove noise during the segmentation process. Diffusion models usually use UNet to model the noise distribution and use denoising-UNet for iterative denoising. However, the U-Net architecture in the traditional diffusion model faces limitations due to its inherent receptive field, making it difficult to capture long-distance related information and solve the problem of global noise distribution. At the same time, the model's ability to capture contextual information is also limited, which leads to low accuracy of image segmentation results. Summary of the invention

[0004] One of the objectives of the present invention is to provide a brain tumor image segmentation method based on a diffusion model to address the deficiencies in the above-mentioned prior art, so as to achieve brain tumor image segmentation more accurately.

[0005] In order to solve the above technical problems, the present invention adopts the following technical solutions: a brain tumor image segmentation method based on a diffusion model, comprising the following steps: establishing a brain tumor image segmentation model, loading a training set and a validation set of images to train the brain tumor image segmentation model to improve its segmentation performance, and then inputting the brain tumor image to be segmented into the trained segmentation model to perform brain tumor image segmentation. The segmentation steps of the brain tumor image segmentation model include:

[0006] (1) Preprocessing of brain tumor images;

[0007] (2) Input the brain tumor image and label into the denoising encoder;

[0008] (3) By introducing a 3D axial attention mechanism module into the denoising encoder, a global long-range dependency model is constructed to improve the model's ability to capture contextual information;

[0009] (4) By introducing a multi-scale feature fusion module in the denoising decoder, the feature extraction capability of the model is enhanced by utilizing multi-scale semantic features, and finally the segmented image is obtained.

[0010] Preferably, the 3D axial attention mechanism module decomposes the 3D attention calculation into three 1D calculations along each dimension to achieve comprehensive 3D modeling of the image, and the module is expressed as:

[0011] Y 1 =LeakyReLU(BN(Conv3D(X)));

[0012] Y Height =H Attn (Y 1 );

[0013] Y Width =W Attn (Y 1 );

[0014] Y Depth =D Attn (Y 1 );

[0015] Y 2 =Y Height +Y Width +Y Depth ;

[0016] Z=LeakyReLU(X+BN(Conv3D(LeakyReLU(Y 2 ))));

[0017] Among them, X represents the input feature map, Y 1 and Y 2 represents the intermediate state feature map, BN and LeakyReLU represent batch normalization and LeakyReLU activation, respectively.

[0018] More preferably, the multi-scale feature fusion module performs a dimensionality reduction operation of convolution on each layer of the five-layer decoder in the model, and then upsamples to the same resolution as the original data, and then fuses the obtained five-layer features and performs global average pooling to extract different weights of features of different scales.

[0019] More preferably, the multi-scale feature fusion module is expressed as:

[0020] Y i =Conv3D(X i),i∈{0,1,2,3,4};

[0021] Z = Concat(upsample(Y i ));

[0022] W = GAP(Z);

[0023] X = Conv3D(W*Z);

[0024] Among them, X i represents the initial feature map of the i-th layer, Y i represents the intermediate state feature map of the i-th layer, X represents the fused output feature map, W represents the fusion weight, Z represents the intermediate state feature map, Conv3D represents a three-dimensional convolutional layer with a kernel size of 1, GAP represents a global average pooling operation, and Concat represents a connection operation.

[0025] In addition, the present invention also provides a brain tumor image segmentation system based on a diffusion model, which comprises:

[0026] The brain tumor image segmentation model is used to load brain tumor images and preprocess them. The brain tumor images and corresponding labels are then input into the denoising encoder. A global long-range dependency model is constructed by introducing a three-dimensional axial attention mechanism module into the denoising encoder to improve the model's ability to capture contextual information. Finally, a multi-scale feature fusion module is introduced into the denoising decoder to enhance the model's feature extraction capability using multi-scale semantic features, and finally a segmented image is obtained.

[0027] A brain tumor image training module, used to train the segmentation model using a sufficient number of brain tumor images to improve its segmentation performance;

[0028] The brain tumor image segmentation module is used to input the brain tumor image to be segmented into the trained segmentation model to perform brain tumor image segmentation.

[0029] The system operates using the diffusion model-based brain tumor image segmentation method described above.

[0030] Compared with the prior art, the present invention has the following beneficial effects:

[0031] 1. By introducing a three-dimensional axial attention mechanism module in the denoising encoder to build global long-range dependencies, the model's ability to capture contextual information is significantly improved, which can effectively solve the limitations of convolution and thus enhance the interpretability of the model.

[0032] 2. By using a multi-scale feature fusion module in the denoising decoder, this module can adaptively fuse features of different scales, thereby effectively utilizing the rich texture and contextual information in the multi-scale features, which helps to enhance the feature extraction capability of the model and thus improve the accuracy of brain tumor image segmentation. BRIEF DESCRIPTION OF THE DRAWINGS

[0033] Figure 1 It is a schematic diagram of the process of the present invention;

[0034] Figure 2 Schematic diagram of the model structure of the method of the present invention;

[0035] Figure 3 Schematic diagram of the denoising network structure of the method of the present invention;

[0036] Figure 4 Schematic diagram of the structure of a three-dimensional axial attention mechanism (AAMM) module in the embodiment;

[0037] Figure 5 Schematic diagram of the structure of the multi-scale feature fusion (MSFFM) module in the embodiment;

[0038] Figure 6 Schematic diagram of the segmentation experiment results in the embodiment. DETAILED DESCRIPTION

[0039] To facilitate understanding by those skilled in the art, the present invention is further described below in conjunction with embodiments and drawings. An existing brain tumor image dataset is used as a training set image of this embodiment, and the contents mentioned in the implementation manner are not intended to limit the present invention.

[0040] like Figure 1 The brain tumor image segmentation method based on the diffusion model is shown in the figure. Figure 2 The multi-scale diffusion model (MSDM) shown in Figure 2 is shown in Figure 2. Figure 3 As shown, compared with the traditional medical image segmentation technology that directly inputs the original image data to predict the final segmentation label, the method of this embodiment learns the denoising process by using the original image and the noise segmentation label as input. The model segmentation network is mainly composed of an auxiliary feature encoder and a denoising network. By introducing a three-dimensional axial self-attention mechanism (AAMM) module in the last layer of the denoising encoder to build a global long-range dependency, the model's ability to capture contextual information is significantly improved. Specifically, this embodiment decomposes the three-dimensional attention calculation into three 1D calculations along each dimension, thereby achieving comprehensive three-dimensional modeling of the image and reducing computational complexity. This method is also conducive to overcoming the limitations of the convolution method, improving the model's contextual information interaction capabilities, and enhancing the learning of deep semantic information.

[0041] By using a multi-scale feature fusion module in the denoising decoder, features of different scales are adaptively fused, the rich texture and context information in the multi-scale features are effectively utilized, and the feature extraction ability of the model is enhanced, thereby improving the accuracy of brain tumor image segmentation.

[0042] In this embodiment, the method combines a diffusion probability model (DDPM) with a fully supervised learning framework to improve the performance of a brain tumor image segmentation model. DDPM generates brain tumor images in a latent space and helps the model extract and retain key tumor feature information through a process of gradually adding and removing noise, while reducing background noise and unnecessary interference. The smooth data distribution in the latent space helps the segmentation model to better focus on the tumor area and enhances the model's ability to learn tumor features. In addition, the data generated by DDPM is highly consistent with the original image in details and boundaries, which enables the model to obtain more refined training signals in fully supervised learning. By comparing the generated data with the real data, the model can more accurately learn the complex structure and change pattern of the tumor, thereby improving the accuracy and reliability of the segmentation results.

[0043] Those skilled in the art should know that the diffusion model uses parameterized Markov chains and variational inference to learn and generate data distribution p(x) by gradually adding and removing noise over multiple time steps. The whole process can be divided into two stages: forward diffusion and reverse diffusion, as shown below:

[0044]

[0045] In this embodiment, if Figure 4 As shown in the figure, the AAMM module decomposes the 3D attention calculation into three 1D calculations along each dimension, thereby achieving comprehensive 3D modeling of the image and reducing computational complexity. This design effectively builds a global long-range dependency model, overcomes the limitations of the convolutional method, and improves the model's ability to capture contextual information. The specific details of the AAMM module are shown in the following formula:

[0046] Y 1 =LeakyReLU(BN(Conv3D(X)));

[0047] Y Height =H Attn (Y 1 );

[0048] Y Width =W Attn (Y 1 );

[0049] Y Depth =D Attn (Y 1 );

[0050] Y 2 =Y Height +Y Width +Y Depth ;

[0051] Z=LeakyReLU(X+BN(Conv3D(LeakyReLU(Y 2 ))));

[0052] Among them, X represents the input feature map, Y 1 and Y 2 represents the intermediate state feature map, BN and LeakyReLU represent batch normalization and LeakyReLU activation, respectively.

[0053] like Figure 5 As shown, for the multi-scale feature fusion module (MSFFM), in order to effectively utilize the rich texture and context information in the multi-scale features, this embodiment uses multi-scale feature fusion in the decoder. Through the five-layer decoder of the 3D-UNet backbone, each layer performs a dimensionality reduction operation of convolution (kernel size is 1×1), and then upsamples to the same resolution as the original data. Then, the obtained five-layer feature fusion is global average pooling, aiming to extract different weights of features of different scales, and finally obtain the fused segmented image. In addition, this embodiment also uses a variety of loss functions and carefully calibrates them to optimize the performance of the network in brain tumor segmentation.

[0054] Y i =Conv3D(X i ),i∈{0,1,2,3,4};

[0055] Z = Concat(upsample(Y i ));

[0056] W = GAP(Z);

[0057] X = Conv3D(W*Z);

[0058]

[0059] Among them, X i represents the initial feature map of the i-th layer, Y i represents the intermediate state feature map of the i-th layer, X represents the fused output feature map, W represents the fusion weight, Z represents the intermediate state feature map, Conv3D represents a three-dimensional convolutional layer with a kernel size of 1, GAP represents a global average pooling operation, and Concat represents a connection operation.

[0060] Figure 6The schematic diagram shows the segmentation results of the existing brain tumor image dataset after using this method to segment the image. Figure 6 In the figure, the column from left to right is the label, the second column is the result obtained by using the segmentation method provided by this embodiment, which is closest to the label, and the third, fourth and fifth columns are the segmentation results of the current mainstream classic model. In comparison, the method proposed in this embodiment obtains better segmentation results.

[0061] The present invention uses a fully supervised learning method to perform brain tumor image segmentation using a large amount of labeled data. This method makes full use of the detailed information in the labeled data, thereby improving the accuracy of the model in the brain tumor segmentation task. Fully supervised learning relies on accurate labeled data to ensure that the model can learn the complex characteristics of brain tumors and show excellent performance in image segmentation. Especially in medical image analysis, accurate segmentation results are crucial for diagnosis and treatment planning. Therefore, brain tumor image segmentation through fully supervised learning not only improves segmentation accuracy, but also provides reliable support for medical decision-making.

[0062] In addition, by utilizing the three-dimensional axial attention mechanism (AAMM) module and the denoising encoder, the model can better understand and utilize global contextual information, thereby improving the segmentation accuracy of complex lesion areas. The multi-scale feature fusion (MSFFM) module adaptively fuses features of different scales in the denoising decoder, effectively utilizing the rich texture and contextual information in the multi-scale features, enhancing the feature extraction capability of the model, thereby improving the accuracy of brain tumor image segmentation. It can be seen that the method of the present invention provides a new research direction and tool for the field of brain tumor image analysis. It is an efficient, accurate solution with good generalization ability, which is of great significance to improving the accuracy and efficiency of clinical diagnosis.

[0063] In order to make it easier for ordinary technicians in the field to understand the improvements of the present invention over the prior art, some drawings and descriptions of the present invention have been simplified, and the above-mentioned embodiments are preferred implementation schemes of the present invention. In addition, the present invention can also be implemented in other ways. Any obvious replacement without departing from the concept of the present technical solution is within the protection scope of the present invention.

Claims

1. A brain tumor image segmentation method based on a diffusion model, characterized by: A brain tumor image segmentation model is established, and training set and validation set images are loaded to train the brain tumor image segmentation model to improve its segmentation performance. The brain tumor image to be segmented is then input into the trained segmentation model for brain tumor image segmentation. The segmentation steps of the brain tumor image segmentation model include: (1) Preprocessing of brain tumor images; (2) Input the brain tumor image and label into the denoising encoder; (3) By introducing a 3D axial attention mechanism module into the denoising encoder, a global long-range dependency model is constructed to improve the model's ability to capture contextual information; (4) By introducing a multi-scale feature fusion module in the denoising decoder, the feature extraction capability of the model is enhanced by utilizing multi-scale semantic features, and finally the segmented image is obtained.

2. The brain tumor image segmentation method based on diffusion model according to claim 1, characterized in that: The 3D axial attention mechanism module decomposes the 3D attention calculation into three 1D calculations along each dimension to achieve comprehensive 3D modeling of the image. The module is expressed as: Y1=LeakyReLU(BN(Conv3D(X))); Y Height =H Attn (Y1); Y Width =W Attn (Y1); Y Depth =D Attn (Y1); Y2=Y Height +Y Width +Y Depth ; Z=LeakyReLU(X+BN(Conv3D(LeakyReLU(Y2)))); Among them, X represents the input feature map, Y1 and Y2 represent the intermediate state feature maps, BN and LeakyReLU represent batch normalization and LeakyReLU activation respectively.

3. The brain tumor image segmentation method based on diffusion model according to claim 1, characterized in that: The multi-scale feature fusion module performs a dimensionality reduction operation by convolution on each layer of the five-layer decoder in the model, and then upsamples to the same resolution as the original data. The obtained five-layer features are then fused and globally averaged pooled to extract different weights of features at different scales.

4. The brain tumor image segmentation method based on diffusion model according to claim 3 is characterized in that: The multi-scale feature fusion module is expressed as: AND i =Conv3D(X i ),i∈{0,1,2,3,4}; Z=Concat(upsample(Y i )); W = GAP(Z); X = Conv3D(W*Z); Among them, X i represents the initial feature map of the i-th layer, Y i represents the intermediate state feature map of the i-th layer, X represents the fused output feature map, W represents the fusion weight, Z represents the intermediate state feature map, Conv3D represents a three-dimensional convolutional layer with a kernel size of 1, GAP represents a global average pooling operation, and Concat represents a connection operation.

5. A brain tumor image segmentation system based on diffusion model, characterized in that: include: The brain tumor image segmentation model is used to load brain tumor images and preprocess them. The brain tumor images and corresponding labels are then input into the denoising encoder. A global long-range dependency model is constructed by introducing a three-dimensional axial attention mechanism module into the denoising encoder to improve the model's ability to capture contextual information. Finally, a multi-scale feature fusion module is introduced into the denoising decoder to enhance the model's feature extraction capability using multi-scale semantic features, and finally a segmented image is obtained. A brain tumor image training module, used to train the segmentation model using a sufficient number of brain tumor images to improve its segmentation performance; The brain tumor image segmentation module is used to input the brain tumor image to be segmented into the trained segmentation model to perform brain tumor image segmentation.

6. The brain tumor image segmentation system based on diffusion model according to claim 5, characterized in that: The method is implemented by the brain tumor image segmentation method based on the diffusion model as claimed in any one of claims 1 to 4.

Images