Three-dimensional brain tumor image segmentation method based on multi-level mamba architecture

Abstract

The application provides a three-dimensional brain tumor image segmentation method based on a multi-level Mamba architecture, which constructs a SegMamba3D network, inputs an original three-dimensional brain tumor MRI image into the network, performs encoding and decoding operations, and outputs a predicted brain tumor segmentation mask. The SegMamba3D network comprises an encoder and a decoder. The encoder performs multi-layer encoding operations on the input sample, performs a downsampling operation at each layer first, and then outputs an encoded feature map through a Mamba Block. The decoder first unifies the size and channel number of the encoded feature maps of different levels output by the encoder through convolution operation, and splices. Then, the spliced feature map is input into a multi-layer perceptron to restore to the size consistent with the input sample. The network is trained to learn a real segmentation mask image, and the trained network is used to output a tumor segmentation result of a three-dimensional brain tumor image.

Description

TECHNICAL FIELD

[0001] The application relates to the field of MRI image segmentation technology, and relates to a three-dimensional brain tumor image segmentation method based on a multilevel Mamba architecture. BACKGROUND

[0002] Computer-aided medical image analysis plays an important role in assisting doctors in clinical diagnosis and determining the best treatment plan, especially medical image segmentation, which can help doctors quickly locate the target area, not only improving the efficiency and accuracy of diagnosis, but also promoting the deep integration of clinical practice and scientific research.

[0003] In the field of medical image segmentation, models based on deep learning have been widely explored. Traditional convolutional neural networks (CNN) use three-dimensional convolution to extract features from three-dimensional images, and use pooling layers, upsampling layers and activation functions to complete the segmentation of three-dimensional images, but may have overfitting and gradient disappearance problems. The Transformer architecture has shown extraordinary ability in modeling global relationships. Architectures based on Vision Transformers (ViTs) can effectively capture global information in images through self-attention mechanisms, thereby improving image segmentation capabilities, making significant progress in 3D medical image segmentation. Correspondingly, it also generates huge computing and memory overheads, especially when dealing with large three-dimensional data, which is prone to memory limitations and computing bottlenecks. SUMMARY

[0004] To overcome the shortcomings of the prior art, the application provides a three-dimensional brain tumor image segmentation method based on a multilevel Mamba architecture, which performs lightweight design on the decoder and introduces an SS3D module to improve the speed and accuracy of three-dimensional medical image segmentation.

[0005] The three-dimensional brain tumor image segmentation method based on the multilevel Mamba architecture comprises the following steps:

[0006] Step 1: Use the original three-dimensional brain tumor MRI image as a sample, and label the real segmentation mask image as a label to construct a training set.

[0007] Step 2: Construct a SegMamba3D network, input the sample obtained in step 1 into the network, and perform encoding and decoding operations to output a predicted brain tumor segmentation mask.

[0008] The SegMamba3D network comprises an encoder and a decoder. The encoder performs a multi-layer encoding operation on the input sample, and at each layer, a downsampling operation is first performed, and then a Mamba Block introducing an SS3D layer is used to output an encoded feature map. The decoder first unifies the size and channel number of the encoded feature maps of different levels output by the encoder through convolution operation, and splices them. Then, the spliced feature map is input into a multi-layer perceptron to restore the size consistent with the input sample, and a predicted brain tumor segmentation mask is obtained.

[0009] The Mamba Block comprises a normalization layer, a linear layer, a convolution layer, an SS3D layer, and a residual connection, so that the network can more easily pass the gradient during back propagation, thereby solving the gradient vanishing problem in deep network training:

[0010] F1 = Linear(LayerNorm(F))

[0011] F2 = SiLU(DWConv(F1))

[0012] F3 = SS3D(F2)

[0013] F4 = Linear(LayerNorm(F3))

[0014] F5 = F4 + F

[0015] F6 = Linear(LayerNorm(F5))

[0016] F7 = Linear(GELU(DWConv(F6)))

[0017] M = F5 + F7

[0018] Wherein, LayerNorm() represents the normalization layer, Linear() represents the linear layer, DWConv() represents the channel-wise convolution layer, SS3D() represents the SS3D layer, SiLU() represents the SiLU activation function, and GELU() represents the GELU activation function. F represents the input feature of the Mamba Block, M represents the output feature of the Mamba Block, and F1-F7 represent the intermediate features of the Mamba Block.

[0019] The multi-layer perceptron performs convolution, normalization and upsampling operations on the spliced feature map:

[0020] Q1 = Conv3d(4C, C)

[0021] Q2 = ReLU(BatchNorm(Q1))

[0022] Q3 = Conv3d(C, N CLS )(Q2)

[0023] Q4 = Upsample(D x H x W)(Q3)

[0024] Wherein, Conv3d() represents three-dimensional convolution, C represents the channel dimension size of the input sample, ReLU() represents the ReLU activation function, BatchNorm() represents batch normalization, N CLS is the class number of the real segmentation mask image, Upsample() is the up sampling operation, D, H and W represent the depth, height and width of the input sample respectively.

[0025] As preferred, the up sampling operation is trilinear interpolation, and interpolation is performed on the image in three dimensions of width, height and depth.

[0026] Step 3, comparing the brain tumor segmentation mask predicted in step 2 with the real segmentation mask image labeled in step 1, calculating the loss value, and adjusting the model parameters according to the loss value.

[0027] Step 4, inputting the three-dimensional brain tumor MRI image into the SegMamba3D network trained in step 3 to obtain the tumor segmentation result in the image.

[0028] The present application has the following beneficial effects:

[0029] By adopting the above technical scheme, the overall network framework is designed by using a classic encoder-decoder, the Mamba is integrated into the encoder, as a state space model (SSM), which models the long-range dependency of sequence data, and has significant storage efficiency and computing speed compared with the Transformer. The hierarchical structure of the encoder can capture multi-level features and improve the performance of semantic segmentation. The lightweight design of the decoder reduces the complexity of the network framework, and to some extent avoids the occurrence of overfitting. BRIEF DESCRIPTION OF DRAWINGS

[0030] Figure 1 The flowchart of the present application.

[0031] Figure 2 The network model framework diagram of the present application.

[0032] Figure 3 The framework diagram of the Mamba module of the present application.

[0033] Figure 4 The path schematic diagram of the SS3D scanning three-dimensional image in the Mamba module of the present application.

[0034] Figure 5A framework diagram of the MLP module of the present application. DETAILED DESCRIPTION

[0035] The present application is further explained with reference to the accompanying drawings, wherein:

[0036] The present application encompasses any alternatives, modifications, equivalent methods and solutions made within the spirit and scope of the present application as defined by the claims. Further, in order to give the public a better understanding of the present application, some specific details are described in the following detailed description of the present application. The present application can also be completely understood without the description of these details by those skilled in the art.

[0037] The present application provides a three-dimensional brain tumor image segmentation method based on a multi-level Mamba architecture, as shown in Figure 1 The method specifically comprises the following steps:

[0038] Step 1: In this embodiment, the BraTs tumor dataset is used as the original data source. After preprocessing, the four modalities of each original three-dimensional brain tumor MRI image, FLAIR, T1w, T1gd and T2w, are combined in the same three-dimensional image, and they are divided into a training set, a validation set and a test set together with the corresponding real segmentation mask image. The data preprocessing steps are as follows:

[0039] s1.1, read the four modalities of FLAIR, T1w, T1gd and T2w images, and separate the data information and spatial conversion information.

[0040] s1.2, the size of the four modalities of the original image is (240, 240, 155), the data information of each dimension of each modality image is normalized and a channel dimension is added, and the size of the final image data information is (1, 240, 240, 155).

[0041] s1.3, create a MetaTensor object containing the final image data information and the corresponding spatial conversion information; then rearrange it to conform to the direction of the RAS coordinate system, and finally separate the data information from the MetaTensor object containing the data and spatial information and convert it to a normal Numpy data.

[0042] s1.4, cut off the area with pixel value of 0, and cut the image from (1, 240, 240, 155) to (1, 128, 128, 128).

[0043] s1.5, splice the four modalities of the image processed by the above steps into a three-dimensional image along the first dimension, and then divide all the processed data into a training set, a validation set and a test set according to a certain proportion.

[0044] Step 2, build a SegMamba3D network as shown in Figure 2 , input the sample obtained in step 1 into the network, and output the predicted brain tumor segmentation mask after encoding and decoding operations. The specific steps are as follows:

[0045] s2.1, the SegMamba3D network performs 4-layer encoding operation on the input sample, first performs downsampling operation at each layer, and then outputs an encoded feature map through two consecutive Mamba blocks.

[0046] For an input sample with a size of DxHxWxC, first perform a downsampling operation through a three-dimensional convolution with a stride of 4, padding of 3, and a convolution kernel size of 7, and then perform feature extraction through two consecutive Mamba blocks to obtain an encoded feature map with a size of Repeat the above steps three times to obtain encoded feature maps with sizes of These features provide high-resolution coarse features and low-resolution fine-grained features, which can improve the performance of semantic segmentation. C1=32, C2=64, C3=160, and C4=256.

[0047] Figure 3 is a structural diagram of the Mamba Block, including a normalization layer, a linear layer, a convolution layer, an SS3D, and an activation function:

[0048] F1=Linear(LayerNorm(F))

[0049] F2=SiLU(DWConv(F1))

[0050] F3=SS3D(F2)

[0051] F4=Linear(LayerNorm(F3))

[0052] F5=F4+F

[0053] F6=Linear(LayerNorm(F5))

[0054] F7=Linear(GELU(DWConv(F6)))

[0055] M=F5+F7

[0056] where LayerNorm() represents a normalization layer, Linear() represents a linear layer, DWConv() represents a channel-wise convolution layer, SS3D() represents an SS3D layer, SiLU() represents a SiLU activation function, and GELU() represents a GELU activation function. F represents the input feature of the Mamba Block, M represents the output feature of the Mamba Block, and F1-F7 represent intermediate features of the Mamba Block.

[0057] In the Mamba Block, the normalization layer adopts Layer Norm, which normalizes all features of each sample, eliminating the size relationship between different samples but preserving the size relationship between different features within a sample. The convolution layer adopts channel-wise convolution, which has lower parameter quantity and operation cost compared to the conventional convolution. The SS3D module specifies the scanning path of the input three-dimensional image, as shown in Figure 4 , the entire three-dimensional image is scanned along the route specified in the figure, making full use of each possible unrolled sequence, thereby providing an in-depth analysis of the entire three-dimensional image and achieving better segmentation results. The activation function uses GELU and SiLU. In addition, the design of the module also utilizes the residual connection idea of ResNet, enabling the network to more easily pass gradients during backpropagation, thereby solving the gradient vanishing problem in deep network training.

[0058] s2.2, M i represents the output feature of the 2i-th Mamba Block. For the four levels of encoded feature maps M i output by the encoder, they are uniformly processed by a multi-layer perceptron layer to have a size of , and then spliced according to the fourth dimension to form a feature map P with a size of :

[0059]

[0060] where i = 1, 2, 3, 4, C i represents the number of channels of the encoded feature map M i , and Concat() is a splicing operation.

[0061] s2.3, as shown in Figure 5 , for the spliced feature map P, a convolution, normalization, and upsampling operation are performed to output a predicted brain tumor segmentation mask with a size of D x H x W x N CLS :

[0062] Q1 = Conv3d(4C, C)

[0063] Q2 = ReLU(BatchNorm(Q1))

[0064] Q3 = Conv3d(C, N CLS )(Q2)

[0065] Q4 = Upsample(D x H x W)(Q3)

[0066] wherein Conv3d() represents a three-dimensional convolution, C represents a channel dimension size of an input sample, ReLU() represents a ReLU activation function, BatchNorm() represents batch normalization, N CLS is a class number of a real segmentation mask image, Upsample() is an up-sampling operation, D, H and W represent depth, height and width of an input sample respectively.

[0067] Step 3, compare the predicted brain tumor segmentation mask output by the SegMamba3D network and the real segmentation mask image corresponding to the sample, calculate the loss value DiceLoss, adjust the model parameters, and then perform segmentation processing on the next image. After all the images are processed, the next round of prediction is performed, and the final prediction model is obtained after all the rounds are completed:

[0068]

[0069] wherein N is the number of pixel points of an input sample, y n is a label, p n is a probability predicted by a model.

[0070] After training for 800 rounds on the BraTs2017 and BraTs2021 data sets respectively, it can be seen that the average Dice score is about 81.74 and 86.87 respectively, which indicates that the average overlap between the predicted image and the real image is 81.74% and 86.87% respectively on the two data sets.

[0071] The embodiments of the present application are described in detail above in combination with the drawings, but the present application is not limited to the described embodiments. For those skilled in the art, various changes, modifications, replacements and variations of these embodiments including components can be made without departing from the principles and spirits of the present application, and still fall within the protection scope of the present application.

Claims

1. A three-dimensional brain tumor image segmentation method based on a multi-level Mamba architecture, characterized by: This method segments the location of a brain tumor from a three-dimensional MRI image containing the tumor, specifically including the following steps: Step 1: Use the original 3D brain tumor MRI image as the sample and the labeled real segmentation mask image as the label to construct a training set; Step 2: Construct a SegMamba3D network including an encoder and a decoder. Input the samples obtained in Step 1 into the network and output the predicted brain tumor segmentation mask. The encoder performs multi-layer encoding operations on the input sample. It first performs a downsampling operation at each layer, and then outputs an encoded feature map by introducing the Mamba Block of the SS3D layer. The decoder first unifies the size and number of channels of the encoded feature maps at different levels output by the encoder through convolution operation and then concatenates them. The concatenated feature map is then input into a multilayer perceptron to restore it to the same size as the input sample, thereby obtaining the predicted brain tumor segmentation mask. The Mamba Block that introduces the SS3D layer includes a normalization layer, a linear layer, a convolutional layer, an SS3D layer, and a residual connection operation: F1 = Linear(LayerNorm(F)) F2 = SiLU(DWConv(F1)) F3 = SS3D(F2) F4 = Linear(LayerNorm(F3)) F5 = F4 + F F6 = Linear(LayerNorm(F5)) F7 = Linear(GELU(DWConv(F6))) M = F5 + F7 Where LayerNorm() represents a normalization layer, Linear() represents a linear layer, DWConv() represents a channel-wise convolutional layer, SS3D() represents an SS3D layer, SiLU() represents the SiLU activation function, GELU() represents the GELU activation function; F represents the input features of the Mamba Block, M represents the output features of the Mamba Block, and F1 to F7 represent the intermediate features of the Mamba Block. Step 3: Compare the brain tumor segmentation mask predicted in Step 2 with the actual segmentation mask image labeled in Step 1, calculate the loss value, and adjust the model parameters based on the loss value. Step 4: Input the 3D brain tumor MRI image into the SegMamba3D network trained in Step 3 to obtain the tumor segmentation result in the image.

2. The three-dimensional brain tumor image segmentation method based on a multi-level Mamba architecture as described in claim 1, characterized in that: For multiple modalities of the original three-dimensional brain tumor MRI image, the data information and spatial transformation information are extracted, the images of multiple modalities are merged into the same three-dimensional image, and the edge regions with a pixel value of 0 are cropped out as samples.

3. The three-dimensional brain tumor image segmentation method based on a multi-level Mamba architecture as described in claim 1, characterized in that: The encoder performs a four-layer encoding operation on the input samples. At each layer, a downsampling operation is performed first, followed by two consecutive Mamba Blocks introduced into the SS3D layer, outputting an encoded feature map M. i .

4. The three-dimensional brain tumor image segmentation method based on a multi-level Mamba architecture as described in claim 1, characterized in that: The multilayer perceptron performs convolution, normalization, and upsampling operations on the stitched feature maps: Q1 = Conv3d(4C,C) Q2 = ReLU(BatchNorm(Q1)) Q3=Conv3d(C,N CLS (Q2) Q4 = Upsample(D×H×W)(Q3) Where Conv3d() represents 3D convolution, C represents the channel dimension of the input sample, ReLU() represents the ReLU activation function, BatchNorm() represents batch normalization, and N CLS The number of categories in the true segmentation mask image is given by Upsample(), which is an upsampling operation. D, H, and W represent the depth, height, and width of the input sample, respectively.

5. The three-dimensional brain tumor image segmentation method based on a multi-level Mamba architecture as described in claim 4, characterized in that: The upsampling operation is a trilinear interpolation, which interpolates the image in the width, height, and depth dimensions respectively.

6. The three-dimensional brain tumor image segmentation method based on a multi-level Mamba architecture as described in claim 1, characterized in that: Use DiceLoss as the loss value: Where N is the number of pixels in the input sample, y n p is the label of the input sample. n Predict the output probability of the SegMamba3D network.

7. A computer-readable storage medium having a computer program stored thereon, which, when executed in a computer, causes the computer to perform the method of any one of claims 1 to 6.

Images