Prostate MRI image segmentation method based on Mamba-Unet
By introducing Mamba-Unet structure and improved edge perception module into the UNet network, and combining attention bridges for feature fusion, UNet solves the problems of overfitting and boundary details loss in prostate MRI image segmentation, achieving higher segmentation accuracy and lower computing requirements.
Patent Information
- Application Number
- CN202411627206.7
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-11-14
- Publication Date
- 2025-03-07
AI Technical Summary
The existing UNet networks have problems such as overfitting, loss of boundary details and poor fusion of different resolution features in prostate MRI image segmentation task.
Using the Mamba-Unet-based prostate MRI image segmentation method, the feature extraction is performed by introducing a 6-layer structure Mamba-Unet network, using a convolution module and a parallel visual Mamba layer (PVM Layer), and boundary information is extracted through the improved edge perception module and edge feature guidance module, combining spatial and channel attention segments to perform multi-scale feature fusion.
The complexity of the model and calculation requirements are reduced, segmentation accuracy is improved, the lesion boundary information in the prostate MRI image is effectively extracted, and the effective fusion of different scale features between the encoder and the decoder is realized.
Smart Images

Figure CN119579627A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the technical field of prostate MRI image segmentation, and in particular to a prostate MRI image segmentation method based on Mamba-Unet. Background Art
[0002] For the prostate MRI image segmentation task, the use of the existing UNet network may have some shortcomings. For example, (1) the prostate MRI image data is small, and the UNet network constructed using full convolution has many parameters, which is prone to overfitting during model training; (2) the UNet network undergoes multiple downsampling during the encoding stage, resulting in loss of detailed information on the lesion boundary, which leads to uneven boundaries in the segmentation results; (3) the jump connection directly fuses the low-resolution information in the encoding with the high-resolution information in the decoding, but these two different resolution information have semantic differences. Direct fusion will lose some spatial information, resulting in poor segmentation results. Summary of the invention
[0003] The purpose of the present invention is to provide a prostate MRI image segmentation method based on Mamba-Unet, which is used to solve the problems raised in the background technology and facilitate promotion.
[0004] In order to achieve the above-mentioned object of the invention, the technical solution adopted by the present invention is:
[0005] A prostate MRI image segmentation method based on Mamba-Unet includes a 6-layer Mamba-Unet, which is mainly composed of an encoder, a decoder and a jump connection. The shallow feature extraction of the first 3 layers uses a convolution module (ConvBlock), each layer includes a standard convolution and a 3×3 convolution kernel and a maximum pooling operation. A novel parallel visual Mamba layer (PVM Layer structure) is introduced in the 4th to 6th layers. The decoder part keeps the same settings as the encoder. In order to extract the edge detail information of the lesion in the MRI image, the output of the 2nd layer and the 6th layer of the encoder are introduced into the improved Parallel Edge-aware Module (PEAM) module to realize the effective extraction of shallow boundary information and deep boundary information. Then, the improved Parallel Edge-guidance Feature Module (PEFM) is introduced to fuse the output of PEAM with the output of each encoder to form a parallel feature fusion structure. The jump connection uses the Spatial Attention Bridge (SAB) module and the Channel Attention Bridge (SAB) module. Bridge, CAB) module performs multi-level and multi-scale information fusion and is connected to the decoder;
[0006] Furthermore, the PVM Layer structure works as follows:
[0007] In order to reduce the computational complexity of the network and capture the relevant information between long sequences, the PVM Layer structure is introduced in the 4th to 6th layers of Mamba-Unet in this paper. Mamba in the PVM Layer is a selective state space model with linear computational complexity. Mamba alleviates the modeling constraints of convolutional neural networks through global receptive fields and dynamic weighting, and provides long-distance modeling capabilities similar to Transformer. The specific operations of the parallel PVM layer are as follows:
[0008] For the output of the 3rd convolution layer from the encoder, the channel C of feature X is passed through the LayerNorm layer and then divided into and Four features, each feature channel is C / 4, and each feature is input into the Mamba structure, and the output is connected through residual connection and adjustment factor to optimize the ability to obtain long-range spatial information. Finally, the four features are combined into a feature Xout with a channel number of C through the concat operation, and output through LayerNorm and Projection operations respectively, as follows:
[0009]
[0010] Out=Pro[LN(X out )] (4)
[0011] Where LN is LayerNorm, Sp is the segmentation operation, Mamba is the Mamba operation, θ is the adjustment factor of the residual connection, Cat is the concat operation, and Pro is the Projection operation. From formula (2), it can be seen that the use of parallel visual Mamba processing features can ensure that the total number of channels processed remains unchanged, maximize parameter reduction, and maintain high accuracy.
[0012] Furthermore, the encoder is followed by an efficient multi-scale attention structure (EMA) and feature aggregation is enriched through cross-spatial learning. EMA's parallel processing of input features avoids the problem of too deep a network depth. For input features, EMA divides the input feature map into multiple sub-feature groups, each group learns different semantics, strengthens the feature learning of the semantic area, and compresses noise at the same time. The parallel subnetwork uses three parallel paths to extract the attention weight descriptors of the grouped feature maps, two of which are 1×1 branches and the third is a 3×3 branch. The 1×1 branch uses a one-dimensional global average pooling operation on the height and width of the image to encode channel information in two spatial directions respectively, and the 3×3 branch captures multi-scale feature representation through 3×3 convolution. The three-way branch parallel structure enables EMA to not only encode cross-channel feature information, but also retain accurate spatial structure information in the channel. EMA does not change the input image size and feature dimension, and has strong generalization.
[0013] Furthermore, the edge perception module works as follows:
[0014] In the PEAM module, the low-level feature f2 and the high-level feature f6 in the encoder are combined to simulate the edge information related to the lesion. The Conv1x1 convolution layer in the EAM module is replaced with the PVM structure. After f2 is upsampled, it is concatenated with f6, and the edge features are learned through three 3×3 convolutions in the convolution module.
[0015] Furthermore, the parallel feature fusion structure works as follows:
[0016] Given input features f2-f6 and boundary features f e , will come from f i The characteristics and f e The upsampled features are element-wise combined with f i Residual connection is performed, and then 3×3 convolution is used to learn fusion features, and the fusion features are input into the PVM structure to build long-distance dependencies. In order to enhance the fusion feature representation, local attention is introduced to explore the local feature channel interaction, and the fusion feature is processed using average pooling (Average Pooling, Avg). The corresponding channel attention weights are obtained through 1D convolution and Sigmoid function and multiplied with the output of PVM to obtain the fusion of each encoder and edge feature as f ei .
[0017] Furthermore, the jump connection working method is as follows:
[0018] The skip connection path uses the SAB module and the CAB module. The combination of SAB and CAB can effectively fuse multi-level features of different scales. The SAB module includes maximum pooling, average pooling, and dilated convolution with shared weights. The CAB module includes global average pooling, concat operation, fully connected layer, and sigmoid activation function.
[0019] Furthermore, in the image segmentation task, the goal is to enable the model to distinguish different regions. The cross entropy loss function helps the model learn how to distinguish different categories more accurately by minimizing the difference between the predicted probability distribution and the true label distribution. The cross entropy loss is used to constrain the model training, which is defined as follows:
[0020]
[0021] Among them, y is the true label, In order to predict the results, in medical image segmentation, we need to face the problem of unbalanced data samples. If the cross entropy loss function is used alone to constrain the model training without considering the category problem, the Dice similarity coefficient is used to measure the overlap between the predicted segmentation and the true segmentation. Combining the cross entropy loss function and the Dice loss function can simultaneously optimize the classification accuracy and segmentation consistency of the model, which is particularly important for medical images that require precise segmentation. The definition is as follows:
[0022]
[0023] Among them, p i g i To predict the point multiplication and addition operation between the segmentation result and the label, the Dice loss expression is:
[0024]
[0025] This paper adopts a combined loss function for constrained model training. The two functions are combined to optimize the loss calculation of the network model locally and overall to improve the segmentation accuracy of the network model. The definition of the combined loss is as follows:
[0026] l=l BCE +l Dice (10)
[0027] Among them, l BCE is the cross entropy loss function, l dice is the Dice loss function.
[0028] As an improvement, the beneficial effects of the present invention are:
[0029] (1) Reduce model complexity and improve segmentation accuracy. The shallow convolution structure of the encoder and the PVMLayer structure can effectively capture the relevant information between long sequences, while reducing network parameters and memory computing requirements.
[0030] (2) Effectively extract the lesion boundary information in prostate MRI images by introducing an improved edge perception module (PEAM) to extract the shallowest and deepest boundary information in the encoder. Then, the extracted boundary information is fused with the output features of each encoder through an improved edge feature guidance module (PEFM), and then connected to the corresponding decoder through a skip connection, thereby providing the decoder with richer boundary detail information.
[0031] (3) Effectively fuse features of different scales between the encoder and decoder. The fused features from different levels in the encoder are fused on the spatial axis using the spatial attention bridge. Then, the global and local information is captured through 1D convolution and full connection in the channel attention bridge. This generates a channel attention map to better integrate channel information and achieve multi-level feature fusion of different scales. BRIEF DESCRIPTION OF THE DRAWINGS
[0032] Figure 1 The network structure of Mamba-Unet of the present invention;
[0033] Figure 2 It is a schematic diagram of the structure of PVM Layer of the present invention;
[0034] Figure 3 This is the structural diagram of the EMA attention mechanism of the present invention;
[0035] Figure 4 This is a schematic diagram of the structure of the edge perception module of the present invention;
[0036] Figure 5 It is a schematic diagram of the structure of the parallel edge guidance feature module of the present invention;
[0037] Figure 6 The test results of different segmentation methods on the PROMISE12 dataset;
[0038] Figure 7 Schematic diagram of the segmentation results of different segmentation algorithms on the PROMISE12 dataset;
[0039] Figure 8 Schematic diagram of ablation experiments on different components on the PROMISE12 dataset;
[0040] Fig. 9 This is a comparison chart of prostate MRI image segmentation methods; DETAILED DESCRIPTION
[0041] In order to make the contents of the present invention more clearly understood, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the accompanying drawings in the embodiments of the present invention.
[0042] This embodiment provides a prostate MRI image segmentation method based on Mamba-Unet. Figure 1 As shown in the figure, Mamba-Unet has a total of 6 layers, which are mainly composed of encoder, decoder and skip connection. The shallow feature extraction of the first 3 layers uses convolutional modules (Conv Block), each layer includes a standard convolution and a 3×3 convolution kernel and a maximum pooling operation. A novel Parallel Vision Mamba Layer (PVMLayer) is introduced in the 4th to 6th layers. The decoder part keeps the same settings as the encoder. In order to extract the edge detail information of the lesion in the MRI image, the output of the 2nd layer and the 6th layer of the encoder are introduced into the improved Parallel Edge-aware Module (PEAM) module to effectively extract the shallow boundary information and the deep boundary information. Then the improved Parallel Edge-guidance Feature Module (PEFM) is introduced to fuse the output of PEAM with the output of each encoder. The skip connection uses the Spatial Attention Bridge (SAB) module and the Channel Attention Bridge (CAB) module to perform multi-level and multi-scale information fusion and connect to the decoder.
[0043] In this embodiment, in order to reduce the computational complexity of the network and capture relevant information between long sequences, the PVM Layer structure is introduced in the 4th to 6th layers of Mamba-Unet in this article. Mamba in the PVM Layer is a selective state space model with linear computational complexity, which performs well in long sequence modeling tasks. Mamba alleviates the modeling constraints of convolutional neural networks through global receptive fields and dynamic weighting, and provides long-distance modeling capabilities similar to Transformer. The specific operations of the parallel PVM layer in this article are as follows:
[0044] For the output of the 3rd convolution layer from the encoder, the channel C of feature X is passed through the LayerNorm layer and then divided into and Four features, each with a feature channel of C / 4, are input into the Mamba structure, and the output is subjected to residual connection and adjustment factors to optimize the ability to obtain long-range spatial information
[21] . Finally, the four features are combined into a feature Xout with a channel number of C through the concat operation, and output through the LayerNorm and Projection operations respectively. The details are as follows:
[0045]
[0046] Out=Pro[LN(X out )] (4)
[0047] Where LN is LayerNorm, Sp is the segmentation operation, Mamba is the Mamba operation, θ is the adjustment factor of the residual connection, Cat is the concat operation, and Pro is the Projection operation. From formula (2), it can be seen that using parallel visual Mamba to process features can ensure that the total number of channels processed remains unchanged, maximize parameter reduction, and maintain high accuracy.
[0048] In this embodiment, in order to effectively learn channel features without compressing convolution channels, this paper introduces a novel efficient multi-scale attention structure (EMA) after each encoder and enriches feature aggregation through cross-spatial learning. EMA's parallel processing of input features avoids the problem of too deep network depth. The structure of the EMA attention mechanism is as follows: Figure 3 As shown. For input features, EMA divides the input feature map into multiple sub-feature groups, each group learns different semantics, strengthens the feature learning of semantic areas, and compresses noise. The parallel subnetwork uses three parallel paths to extract the attention weight descriptors of the grouped feature maps, two of which are 1×1 branches and the third is a 3×3 branch. The 1×1 branch uses a one-dimensional global average pooling operation on the height and width of the image to encode channel information in two spatial directions respectively. The 3×3 branch captures multi-scale feature representations through 3×3 convolutions. The three-way branch parallel structure enables EMA to not only encode cross-channel feature information, but also retain accurate spatial structure information in the channel. EMA does not change the input image size and feature dimension, and has strong generalization.
[0049] In this embodiment, high-level semantic or position information is needed to promote the edge features of the region of interest. In the PEAM module, the low-level features f2 and high-level features f6 in the encoder are combined to model the edge information related to the lesion.
[0050] like Figure 4Specifically, in order to make the model lighter and reduce the amount of channel calculation while learning more boundary knowledge, this paper replaces the Conv1x1 convolution layer in the EAM module with the PVM structure. Then f2 is upsampled and concatenated with f6. Next, the edge features are learned through three 3×3 convolutions in the Block module.
[0051] In this embodiment, the purpose of introducing an improved edge-guidance feature module (EFM) is to inject edge information related to the boundary into the feature representation to enhance the feature representation of the semantics of the lesion structure. Different feature channels usually contain different semantic information. Therefore, in order to achieve a good integration of shallow edge information and deep information in the encoder, a local channel attention mechanism is introduced to realize the interaction between channels and mine the key information between channels.
[0052] like Figure 5 As shown, given the input features f2-f6 and the boundary feature f e . Will come from f i The characteristics and f e The upsampled features are added element by element and combined with f i Perform residual connection. Then use 3×3 convolution to learn fusion features, and input the fusion features into the PVM structure to build long-distance dependencies. In order to enhance the fusion feature representation, local attention is introduced to explore local feature channel interactions. Use average pooling (Average Pooling, Avg) to separate the fusion features, and then obtain the corresponding channel attention weights through 1D convolution and Sigmoid function and multiply them with the output of PVM. The fusion of each encoder and edge feature is obtained as f ei .
[0053] In this embodiment, the skip connection path uses the SAB module and the CAB module. SAB and CAB have been proven to be effective in improving the convergence ability of the model and enhancing the sensitivity to lesions. The combined use of SAB and CAB can effectively fuse multi-level features of different scales. The SAB module includes maximum pooling, average pooling, and dilated convolution with shared weights. The CAB module includes global average pooling, concat operation, fully connected layer, and sigmoid activation function.
[0054] In this embodiment, in the image segmentation task, the goal is to enable the model to distinguish different regions. The cross entropy loss function helps the model learn how to distinguish different categories more accurately by minimizing the difference between the predicted probability distribution and the true label distribution. Therefore, this paper adopts the cross entropy loss constraint model training, which is defined as follows:
[0055]
[0056] Among them, y is the true label, is the prediction result. In medical image segmentation, we need to face the problem of unbalanced data samples. If the cross entropy loss function is used alone to constrain the model training without considering the category problem, the Dice similarity coefficient is used to measure the overlap between the predicted segmentation and the true segmentation. Combining the cross entropy loss function and the Dice loss function can simultaneously optimize the classification accuracy and segmentation consistency of the model, which is particularly important for medical images that require precise segmentation. The definition is as follows:
[0057]
[0058] Among them, p i g i To predict the point multiplication and addition operation between the segmentation result and the label, the Dice loss expression is:
[0059]
[0060] This paper adopts a combined loss function for constrained model training, combining two functions to optimize the loss calculation of the network model locally and overall to improve the segmentation accuracy of the network model. The definition of the combined loss is as follows:
[0061] l=l BCE +l Dice (10)
[0062] Among them, l BCE is the cross entropy loss function, l dice is the Dice loss function.
[0063] This embodiment provides experimental results and analysis, and the specific implementation method is as follows:
[0064] To verify the proposed method, the experimental data set uses the PROMISE12 data set proposed by the Medical Image Computing and Computer Assisted Intervention Society (MICCAI) in the prostate MRI image segmentation challenge in 2012. It contains 100 prostate MRI data, 1473 training images, and 473 test images. All data go through the same preprocessing process, as follows:
[0065] First, the nii format data is converted to png format; in order to enhance the detail information of small areas in the image and enhance the contrast of the area of interest in the entire image, this paper adopts adaptive histogram equalization processing. In order to improve the clarity of the image and reduce noise, Gaussian filtering is used to further process the image. Then, according to the doctor's manual segmentation result file, the center of the mask is determined, and the center coordinates are used as the origin. The mask boundary is made into a rectangle and expanded by 5 pixels on the upper and lower and left and right boundaries to ensure that the lesion is surrounded. Secondly, the obtained rectangle is applied to the original image, and the corresponding area is intercepted to obtain the lesion area of prostate cancer in the MRI image. Finally, in order to meet the requirements of the input deep neural network model, the lesion area is normalized to 224×224 pixels.
[0066] The experimental environment of this article is: Linux version 6.1.85 system, Nvidia L4 graphics card, the network uses Pytorch 1.13+cu117 deep learning framework for network construction. During the training process, the network parameters are randomly initialized, the batch is configured to 16, the epoch is 50, and the AdamW optimizer is used to update the parameters, lr=0.001, betas=(0.9,0.999), eps=1e-8, weight_decay=1e-2 optimizer: CosineAnnealingLR.
[0067] In order to evaluate the segmentation performance of different segmentation models, this paper uses several evaluation indicators such as Dice similarity coefficient (DSC), intersection over Union (IoU), accuracy, specificity, sensitivity and HD95 to objectively and quantitatively evaluate the segmentation algorithm. The calculation formulas of DSC, IoU, Accuracy, Specificity, Sensitivity and HD95 are as follows:
[0068]
[0069] Among them, TP and FN represent the pixels correctly classified as prostate areas and pixels misclassified as irrelevant background tissue areas in the experiment, TN and FP represent the voxels correctly classified as background tissues and misclassified as prostate areas. HD measures the distance between the algorithm segmentation boundary and the true boundary. and Y are the contour set of the predicted result graph and the contour set of the gold standard segmentation result graph, respectively. Represents the boundary pixels of the predicted segmentation result map The Euclidean distance from the boundary pixel point y of the gold standard segmentation result image.
[0070] In order to further verify the segmentation method proposed in this paper, qualitative and quantitative analysis were performed with the classic UNet, SegNet, ResUNet, CENet, TransUNet, SwinUNet, VM-UNet, and UltraLight-VM-UNet segmentation algorithms.
[0071] like Figure 6 As shown in the figure, compared with other segmentation algorithms, ResUNet, CENet, SegNet and UNet based on full convolution perform poorly in the prostate segmentation task; while TransUNet and SwinUNet based on the combination of Transformer and CNN have better segmentation performance in the prostate segmentation task than the segmentation model based on full convolution. This is because TransUNet and SwinUNet respectively take advantage of the advantages of extracting local features by convolution and extracting global features by Transformer during the modeling process. In contrast, VM-UNet, UltraLight-VM-UNet and Our segmentation models based on Mamba perform best in the prostate segmentation task. Taking DSC, IoU, Accuracy, Specificity, Sensitivity and HD95 as evaluation indicators, our DSC and IoU reach 92.75 and 86.81 respectively, which are 11.62 and 17.33 higher than the UNet model respectively. HD95 is 2.03, which is 11.26 lower than the UNet model. Compared with UNet, the parameter of this method is reduced by 30,900,791. It can be seen that the Mamba-Unet method proposed in this paper is not only lightweight, but also can improve the segmentation accuracy of the prostate.
[0072] The visualization of prostate segmentation results of each segmentation model is as follows Figure 7 As shown in the figure, the first and second columns are prostate MRI images and corresponding labels. Figure 7It can be seen that when segmenting prostate MRI images, the ResUNet segmentation model has the worst segmentation results, with many mis-segmentations and even the lesion area not being segmented. The reason may be that there is less data, resulting in overfitting of the model training. The UNet, SegNet, and CENet segmentation models can segment the approximate area of the lesion, but there are also some under-segmented areas, and the segmentation boundaries are fuzzy and there are many burrs. The main reason is that the segmentation model based on CNN modeling can only extract local features of the image and lose the global information of the image. In addition, in the process of skip connection, the results of multiple downsampling in the encoder are directly fused with the decoder, and there is a large semantic gap between the features from different resolutions, which makes the final segmentation result of the decoder have more burr edges. Compared with the segmentation model built based on CNN, the TransUNet and SwinUNet segmentation models built based on Transformer and CNN have certain advantages. Although they combine the advantages of local modeling and long-distance modeling, from the segmentation visualization results, there are still boundary segmentation models and under-segmentation and mis-segmentation parts. The reason is that in order to obtain ideal segmentation results, the Transformer-based segmentation model needs to experiment on a large amount of data. However, for the task of segmenting medical images, due to the small number of data samples and the large differences in the structure of the data itself, the segmentation results are not ideal.
[0073] Compared with the segmentation models based on CNN and Transformer, the segmentation model based on Mamba has great advantages. This model combines the advantages of CNN and Transformer, and at the same time, the model parameters are low, the model reasoning is fast, and it takes up less computer resources. From the segmentation results of VM-UNet and UltraLight-VM-UNet models, it can be seen that the lesion area in the prostate image can be effectively segmented. However, the boundary of the segmentation result of VM-UNet is still not smooth, while the segmentation result of UltraLight-VM-UNet is relatively smooth, but the phenomenon of over-segmentation and under-segmentation also occurs, which may be caused by the small number of model training rounds and the small number of data samples.
[0074] from Figure 7It can be seen that the Mamba-Unet segmentation model proposed in this paper has the best segmentation results. A novel Mamba structure is introduced into the Unet model, which has long modeling capabilities and low parameters, and is combined with CNN to construct the encoder and decoder. During the construction process, the importance of image lesion edge information to the segmentation result is considered, so the PEAM structure is used between the shallowest and deepest layers of the encoder to extract edge information at different resolutions in different encoders, and then the edge information is fused with the output of each encoder using the PEFM structure to obtain features with edge information in each encoder. In order to alleviate the semantic differences between the encoder and the decoder, spatial attention and channel attention are used between the encoder and the decoder for multi-scale fusion, thereby improving the segmentation accuracy of the model. It can be seen that the Mamba-Unet model proposed in this paper has great application potential in prostate segmentation tasks.
[0075] The model components in this paper are mainly divided into the boundary perception module (PEAM) introduced in the encoder and the improved edge feature guidance module (PEFM). In order to verify the improvement strategy of the segmentation algorithm in this paper, the effects of the above components on the performance of the segmentation algorithm are verified respectively. The experimental results are as follows: Figure 8 As shown:
[0076] Depend on Figure 8 It can be seen that after the encoder in the Baseline network introduces the boundary perception module (EAM) and the edge feature guidance module (EFM), the segmentation effect of the model is improved compared with the Baseline. This shows that the shallow features and deep features extracted by the EAM module have a promoting effect on the model segmentation of the prostate after being fused by the EFM module. When the 1×1 convolution in the EAM module and the EFM module are replaced by the PVM Layer, the segmentation performance of the model is slightly improved, indicating that the PVMLayer structure has certain advantages in long-distance modeling. When the PVM Layer structure is introduced into the EAM module and the EFM module at the same time, the segmentation performance of the model is greatly improved. Compared with the Baseline, DSC, IoU, Accuracy, Specificity and Sensitivity are increased by 1.43, 2.36, 0.75, 0.39 and 1.1 respectively, and HD95 is reduced by 0.91; compared with the Baseline+EAM+EFM model, DSC, IoU, Accuracy, Specificity and Sensitivity are increased by 0.81, 1.21, 0.30, 0.07 and 0.86 respectively, and HD95 is reduced by 0.48.
[0077] In order to prove the segmentation performance of the segmentation algorithm in this paper, it is compared with the prostate MRI image segmentation methods proposed by scholars in recent years. The algorithms are all verified using the PROMISE12 dataset. Fig. 9 As shown in the figure, the highest DSC, IoU, Accuracy and Sensitivity of prostate MRI image segmentation algorithms in recent years are 90.30, 79.24, 99.32 and 94.44 respectively; for the HD95 evaluation index, the best is 3.192. The DSC, IoU, Accuracy, Specificity and Sensitivity of the segmentation algorithm in this paper are 92.75, 86.81, 96.57, 97.74 and 93.30 respectively, and the HD95 is 2.03. Except for the Accuracy and Sensitivity indicators, the segmentation algorithm in this paper is the highest in other evaluation indicators. By comparison, it can be seen that the segmentation algorithm in this paper has better segmentation performance.
[0078] In view of the fuzzy boundary and low contrast of lesions in prostate MRI images, the traditional UNet segmentation network cannot effectively extract boundary information, and the boundary detail information is lost after multiple downsampling, and the problem of feature fusion of different resolutions in the encoder and decoder is solved. A lightweight Mamba-Unet segmentation method is proposed. The PVM Layer structure is introduced in the deep layer of the encoder. As a linear state selection model, the PVM Layer can not only reduce the number of model parameters, but also has good long-distance modeling ability. In order to effectively extract the boundary detail information of lesions in different levels of the encoder, an improved PEAM module is introduced between the shallowest and deepest layers of the encoder, and the extracted boundary detail information is fused with the output of each encoder using the improved PEFM module to increase the detail information of different levels of the encoder. In the jump connection, the spatial attention bridge and channel attention bridge are used to connect the fused output features of the encoder with the corresponding decoder, so as to realize the multi-scale feature fusion between different levels of the encoder and decoder, so as to make the segmentation result smoother. Experimental results show that the segmentation model proposed in this paper can not only have high segmentation accuracy, but also have small parameters and low computational memory. The next step of research will be to continue to expand the validation data and add other data sets to verify the generalization performance of the segmentation method in this paper.
[0079] The above are only preferred embodiments of the present invention and are not intended to limit the present invention. Any modifications, equivalent substitutions and improvements made within the spirit and principles of the present invention shall be included in the protection scope of the present invention.
Claims
1. A prostate MRI image segmentation method based on Mamba-Unet includes a 6-layer Mamba-Unet, which is mainly composed of an encoder, a decoder and a jump connection. The shallow feature extraction of the first three layers uses a convolution module (Conv Block), and each layer includes a standard convolution and a 3×3 convolution kernel and a maximum pooling operation. A novel parallel visual Mamba layer (PVM Layer structure) is introduced in the 4th to 6th layers. The decoder part keeps the same settings as the encoder. In order to extract the edge detail information of the lesion in the MRI image, the output of the 2nd layer and the 6th layer of the encoder are introduced into the improved Parallel Edge-aware Module (PEAM) module to realize the effective extraction of shallow boundary information and deep boundary information. Then, an improved Parallel Edge-guidance Feature Module (PEFM) is introduced to fuse the output of PEAM with the output of each encoder to form a parallel feature fusion structure. The jump connection uses the Spatial Attention Bridge (SAB) module and the Channel Attention Bridge (SAB) module. The AttentionBridge (CAB) module performs multi-level and multi-scale information fusion and is connected to the decoder.
2. The prostate MRI image segmentation method based on Mamba-Unet according to claim 1, characterized in that: The PVM Layer structure works as follows: In order to reduce the computational complexity of the network and capture the relevant information between long sequences, the PVM Layer structure is introduced in the 4th to 6th layers of Mamba-Unet in this paper. Mamba in the PVM Layer is a selective state space model with linear computational complexity. Mamba alleviates the modeling constraints of convolutional neural networks through global receptive fields and dynamic weighting, and provides long-distance modeling capabilities similar to Transformer. The specific operations of the parallel PVM layer are as follows: For the output of the 3rd convolution layer from the encoder, the channel C of feature X is passed through the LayerNorm layer and then divided into and Four features, each feature channel is C4, and each feature is input into the Mamba structure, and the output is connected through residual connection and adjustment factor to optimize the ability to obtain long-range spatial information. Finally, the four features are combined into the feature Xout with the number of channels C through the concat operation, and output through the LayerNorm and Projection operations respectively, as follows: Out=Pro[LN(X out )] (4) Where LN is LayerNorm, Sp is the segmentation operation, Mamba is the Mamba operation, θ is the adjustment factor of the residual connection, Cat is the concat operation, and Pro is the Projection operation. From formula (2), it can be seen that the use of parallel visual Mamba processing features can ensure that the total number of channels processed remains unchanged, maximize parameter reduction, and maintain high accuracy.
3. The prostate MRI image segmentation method based on Mamba-Unet according to claim 1, characterized in that: The encoder is followed by an efficient multi-scale attention structure (EMA) and feature aggregation is enriched through cross-spatial learning. EMA's parallel processing of input features avoids the problem of too deep a network depth. For input features, EMA divides the input feature map into multiple sub-feature groups, each group learns different semantics, strengthens the feature learning of the semantic area, and compresses noise at the same time. The parallel subnetwork uses three parallel paths to extract the attention weight descriptors of the grouped feature maps, two of which are 1×1 branches and the third is a 3×3 branch. The 1×1 branch uses a one-dimensional global average pooling operation on the height and width of the image to encode channel information in two spatial directions respectively, and the 3×3 branch captures multi-scale feature representation through 3×3 convolution. The three-way branch parallel structure enables EMA to not only encode cross-channel feature information, but also retain accurate spatial structure information in the channel. EMA does not change the input image size and feature dimension, and has strong generalization.
4. The prostate MRI image segmentation method based on Mamba-Unet according to claim 3, characterized in that: The edge perception module works as follows: In the PEAM module, the low-level feature f2 and the high-level feature f6 in the encoder are combined to simulate the edge information related to the lesion. The Conv1x1 convolution layer in the EAM module is replaced with the PVM structure. After f2 is upsampled, it is concatenated with f6, and the edge features are learned through three 3×3 convolutions in the convolution module.
5. The prostate MRI image segmentation method based on Mamba-Unet according to claim 4, characterized in that: The parallel feature fusion structure works as follows: Given input features f2-f6 and boundary features f e , will come from f i The characteristics and f e The upsampled features are element-wise combined with f i Residual connection is performed, and then 3×3 convolution is used to learn fusion features, and the fusion features are input into the PVM structure to build long-distance dependencies. In order to enhance the fusion feature representation, local attention is introduced to explore the local feature channel interaction, and the fusion feature is processed using average pooling (Average Pooling, Avg). The corresponding channel attention weights are obtained through 1D convolution and Sigmoid function and multiplied with the output of PVM to obtain the fusion of each encoder and edge feature as f ei .
6. The prostate MRI image segmentation method based on Mamba-Unet according to claim 1, characterized in that: The jump connection working method is as follows: The skip connection path uses the SAB module and the CAB module. The combination of SAB and CAB can effectively fuse multi-level features of different scales. The SAB module includes maximum pooling, average pooling, and dilated convolution with shared weights. The CAB module includes global average pooling, concat operation, fully connected layer, and sigmoid activation function.
7. The prostate MRI image segmentation method based on Mamba-Unet according to claim 1, characterized in that: In the image segmentation task, the goal is to enable the model to distinguish different regions. The cross entropy loss function helps the model learn how to distinguish different categories more accurately by minimizing the difference between the predicted probability distribution and the true label distribution. The cross entropy loss constraint model training is used, which is defined as follows: Among them, y is the true label, In order to predict the results, in medical image segmentation, we need to face the problem of unbalanced data samples. If the cross entropy loss function is used alone to constrain the model training without considering the category problem, the Dice similarity coefficient is used to measure the overlap between the predicted segmentation and the true segmentation. Combining the cross entropy loss function and the Dice loss function can simultaneously optimize the classification accuracy and segmentation consistency of the model, which is particularly important for medical images that require precise segmentation. The definition is as follows: Among them, p i g i To predict the point multiplication and addition operation between the segmentation result and the label, the Dice loss expression is: This paper adopts a combined loss function for constrained model training. The two functions are combined to optimize the loss calculation of the network model locally and overall to improve the segmentation accuracy of the network model. The definition of the combined loss is as follows: l=l BCE +l Dice (10) Among them, l BCE is the cross entropy loss function, l dice is the Dice loss function.
Citation Information
Cited By
Construction method of bidirectional adaptive reweighted neural network for image processing
CN119903875A
Remote sensing image target detection method based on DMSE-YOLO network
CN119992276A
Vision-Mamba CNN-based end-to-end image compression method and device
CN120088349A
Intensive prediction multi-task learning method combining mixed expert and Mangbar model
CN120182791A
A Dense Prediction Multi-Task Learning Method Combining Mixture of Experts and Mamba Model
CN120182791B