A 2D medical image segmentation method and system based on Mamba and UNet

Abstract

This invention discloses a 2D medical image segmentation method and system based on Mamba and UNet, comprising: collecting and preprocessing a medical image segmentation dataset to obtain a training set; constructing a 2D medical image segmentation model based on Mamba and UNet, wherein the 2D medical image segmentation model includes a block embedding layer, an encoder, a decoder, and a prediction generation layer; designing an adaptive hierarchical loss function based on gradient statistics and training the 2D medical image segmentation model on the training set; and inputting the medical image to be segmented into the trained model to complete the image segmentation. This invention, through the innovative construction of a 2D medical image segmentation model based on Mamba and UNet, can achieve automated and intelligent segmentation of medical images with high segmentation accuracy and efficiency.

Description

Technical Field

[0001] This invention relates to the fields of computer vision deep learning and medical image processing technology, and in particular to a 2D medical image segmentation method and system based on Mamba and UNet. Background Technology

[0002] Medical image segmentation is a key task in computer vision, enabling the precise division of anatomical structures and pathological regions to improve diagnosis, treatment planning, and surgical navigation. Unlike conventional images, medical images present unique challenges such as low contrast, blurred boundaries, inter-class similarity, and noise artifacts, making accurate segmentation particularly difficult.

[0003] In existing technologies, convolutional neural networks (CNNs), particularly UNet, have become the foundational architecture for medical image segmentation. However, the inherent limitations of CNNs in modeling long-range dependencies restrict their performance on complex anatomical structures. Transformers largely address this limitation through their self-attention mechanism, but their quadratic computational complexity and difficulties in preserving the fine local features required for accurate boundary delineation remain problems.

[0004] With the development of state-space models (SSMs), especially Mamba, they have become a promising alternative due to their linear computational complexity and efficient long-range modeling capabilities.

[0005] For example, Chinese patent document CN118447244A discloses a brain tumor image segmentation method based on multi-scale convolution and Mamba structure. It reduces computation by introducing a state-space model and improves the recognition of detailed features by introducing a cascaded residual multi-scale convolution module. Chinese patent document CN120070891A discloses a tumor segmentation method based on Mamba-guided multi-encoder fusion, including: 1. Constructing a multi-encoder segmentation model fusing Mamba and convolutional neural networks (CNN); 2. Training the model and optimizing its parameters on liver and lung cancer CT image datasets; 3. Using the trained model to quickly locate and accurately segment the input CT image, thereby obtaining a binary image of the tumor segmentation result.

[0006] However, existing Mamba-based 2D medical image segmentation methods still have the following problems:

[0007] 1. Existing Mamba-based 2D medical image segmentation methods are generally not ideal for segmenting fine structures such as small lesions or complex organ boundaries.

[0008] 2. Existing Mamba-based 2D medical image segmentation methods with multiple encoders and decoders often suffer from deficiencies in cross-scale feature fusion between encoders and decoders, lacking feedback guidance from high-level semantics to low-level features.

[0009] 3. Existing Mamba-based 2D medical image segmentation methods, if using multi-level loss combinations, typically employ fixed weights or simple learnable weights, which cannot adapt to the dynamic training requirements of different decoder levels. Summary of the Invention

[0010] To address the shortcomings of existing medical image segmentation methods, this invention provides a 2D medical image segmentation method and system based on Mamba and UNet. Through an innovative architecture design of a local-global fusion visual Mamba UNet network (LGFVM-UNet), it can achieve automated and intelligent segmentation of medical images with high segmentation accuracy and efficiency.

[0011] A 2D medical image segmentation method based on Mamba and UNet includes the following steps:

[0012] (1) Collect medical image segmentation datasets and preprocess them to obtain training sets;

[0013] (2) Construct a 2D medical image segmentation model based on Mamba and UNet. The 2D medical image segmentation model includes a patch embedding layer, an encoder, a decoder and a prediction generation layer.

[0014] The block embedding layer is used to map the input image into a semantic feature vector and then input it into the encoder; specifically, the input medical image is passed through a convolutional layer with a 4×4 kernel.

[0015] The encoder consists of multiple stages. In the first stage, the feature map is extracted from the output of the block embedding layer by the Local-Global Fusion Selective State Space 2D Module (LGF-VSS). In each stage except the first stage, the feature map extracted by the LGF-VSS module is received from the output of the previous stage. In each stage except the last stage, the feature map extracted by the LGF-VSS module is downsampled by the corresponding block merging layer before being input into the next stage. In the last stage, the feature map extracted by the LGF-VSS module is directly input into the decoder without going through the block merging layer.

[0016] The decoder consists of multiple stages. The first stage receives the output of the last stage of the encoder through the LGF-VSS module for feature refinement. Each stage except the first stage receives the output of the previous stage through the LGF-VSS module for feature refinement. Each stage uses a multi-level cross-scale feature fusion block (MCFB) to fuse the refined feature map with the feature maps from all stages of the encoder to generate a feature map that contains both high-level semantic understanding and precise spatial positioning information. The feature maps generated in each stage except the last stage are input into the next stage, and the feature map generated in the last stage is input into the prediction generation layer.

[0017] The prediction generation layer upsamples and performs pixel-level classification on the feature map output from the last stage of the decoder to obtain pixel-level mask segmentation results.

[0018] (3) Design an adaptive hierarchical loss function based on gradient statistics to train the 2D medical image segmentation model on the training set;

[0019] (4) Input the segmented medical image into the trained model to complete the image segmentation.

[0020] In this invention, step (2) involves four stages for both the encoder and the decoder.

[0021] In step (2), the LGF-VSS module efficiently fuses local detail features with global contextual information by combining a selective state-space model and multi-scale convolution operations, thereby enhancing the model's feature representation capabilities in visual tasks. Specifically, the LGF-VSS module includes the LGF-SS2D module and a dynamic gating fusion block;

[0022] The LGF-SS2D module captures the global context through four-directional state space scanning, obtains state space features, and extracts multi-scale local features by combining multi-scale convolutions with kernel sizes of 1, 3, and 5.

[0023] The dynamic gated fusion block generates four sets of spatial adaptive fusion weights, which are then weighted and fused with the state space features and multi-scale local features output by the LGF-SS2D module to obtain the fused features.

[0024] The fused features are combined with the original input of the LGF-SS2D module through a dual residual connection mechanism, and the output feature map is normalized by layers.

[0025] Furthermore, the detailed implementation process of the LGF-VSS module is as follows:

[0026] The input feature map is normalized by applying LayerNorm and the feature dimension order is adjusted from (B,H,W,C) to (B,C,H,W). Then, convolution is used to perform preliminary spatial feature extraction on the input features, and a nonlinear representation is introduced through the SiLU activation function to obtain the first-generation feature map in LGF-VSS.

[0027] Inputting the initial feature map from LGF-VSS into LGF-SS2D yields three convolutional branch feature maps and one state space feature map;

[0028] The initial feature map in LGF-VSS is input into the QuadGate dynamic gated fusion block to obtain four sets of spatially adaptive fusion weights;

[0029] The state space feature map and the feature maps of the three convolutional branches are weighted and summed to achieve adaptive fusion of local and global features.

[0030] The fused features are combined with the original LGF-VSS input through a residual connection. The output features are then normalized using a layer and a gating mechanism is employed, whereby multiplication with the original input controls the information flow, enhancing feature selectivity. Finally, the dimensions of the processed feature map are restored from (B,C,H,W) to (B,H,W,C), maintaining the same dimensional order as the input.

[0031] The LGF-SS2D module captures the global context through a four-directional state space scan, as follows:

[0032] The feature map is unfolded into a sequence along four dimensions: forward height, backward height, forward width, and backward width.

[0033] The Mamba model was used to model long-range dependencies of sequences in each direction.

[0034] By fusing the output features from the four directions, we obtain the global context features, i.e., the state space features.

[0035] Furthermore, the detailed implementation steps of LGF-SS2D are as follows:

[0036] The SS2D module is used to perform a four-directional selective state-space scan of the initial feature map in LGF-VSS, which involves unfolding the feature map into a sequence along four dimensions: forward height, backward height, forward width, and backward width. Then, a Mamba model is used to model the long-distance dependencies of the sequence in each direction. Finally, the output features from the four directions are fused to capture the global spatial context information of the feature map, resulting in a state-space feature map.

[0037] While capturing global spatial context information of the feature map using the SS2D module and the Mamba model, three convolutional branches—1×1, 3×3, and 5×5—are used in parallel to extract local features with different receptive fields from the initial feature map within the LGF-VSS. After each branch, the output and input layers obtained from the branch are normalized (LayerNorm) and SiLU activated to enhance feature representation capabilities, resulting in three convolutional branch feature maps.

[0038] The specific process of generating four sets of spatial adaptive fusion weights by the dynamic gated fusion block is as follows:

[0039] Global average pooling is performed on the features input to the LGF-VSS module to obtain channel statistics;

[0040] Channel statistics are used to generate four sets of channel attention weights through a linear layer, a SiLU activation function, and a 1×1 convolution.

[0041] The weights are normalized using the Softmax function to obtain four sets of spatially adaptive fusion weights.

[0042] The working process of the block merging layer is as follows:

[0043] The input feature map is divided into 2×2 local regions, and pixels are reorganized to the channel dimension to achieve a 4-fold channel expansion.

[0044] The number of channels is compressed to twice the size using 1×1 convolution;

[0045] The application layer normalizes the output downsampled feature map.

[0046] The MCFB module fuses the refined feature maps with feature maps from all stages of the encoder. The specific process is as follows:

[0047] The resolution of feature maps from all stages of the encoder output can be unified by upsampling or downsampling;

[0048] Generate spatial-channel dual attention weights and weight features from all stages of the encoder;

[0049] The weighted encoder features are concatenated with the decoder features of the current stage, fused by LGF-VSS, and then output through residual connection.

[0050] Furthermore, the detailed implementation steps of the MCFB module are as follows:

[0051] For each encoder feature map except the feature map corresponding to the current decoder layer, an attention gating mechanism is used to interact with the encoder feature map corresponding to the current decoder layer to generate spatial and channel dual attention weights, and the original feature map is weighted to highlight important information.

[0052] By performing upsampling or downsampling operations, the spatial resolution of all encoder feature maps, including those of the current layer, is uniformly adjusted to the target resolution of the current decoder layer, while the number of channels is adjusted to one-quarter of the target number of channels.

[0053] All feature maps that have undergone resolution and channel adjustments are stitched together along the channel dimension to form a multi-scale fused feature map, and then deep feature map fusion and enhancement are performed using LGF-VSS.

[0054] The fused feature map is residually connected to the current decoder layer feature map, and finally layer normalization (LayerNorm) is applied to stabilize the feature distribution, outputting the enhanced final feature map of the current decoder stage.

[0055] The implementation steps of the prediction generation layer are as follows:

[0056] A linear transformation is applied to the low-resolution, high-channel-count feature map output by the decoder to expand the number of channels of the input features to four times the original number of channels, thus preparing sufficient information for subsequent spatial upsampling.

[0057] Layer normalization and tensor reshaping operations are performed on the expanded feature map. The normalized high-resolution feature map is then input into a 1×1 convolutional layer, mapping the number of channels in the feature map to the number of classes in the target segmentation task. Next, a Softmax function is applied to the class score vector at each pixel position output by the 1×1 convolution, converting the value at each pixel position into a probability distribution representing the probability that the pixel belongs to each class. The class index with the highest probability is taken as the final predicted class for that pixel, thus generating a complete pixel-level semantic segmentation mask.

[0058] In step (3), an adaptive hierarchical loss function based on gradient statistics is designed, specifically including:

[0059] Calculate the gradient magnitude of the decoder output at each stage except the last stage:

[0060]

[0061] Where |θ| represents the total number of network parameters, It is the loss of the decoder in the l-th stage excluding the last stage;

[0062] The loss weights for each stage are updated using momentum smoothing:

[0063]

[0064] Where α∈[0,1] is the momentum coefficient controlling the retention of historical information, exp represents the exponential function, and τ>0 is used as a temperature parameter to control the sharpness of the weight distribution; if If it does not exist, meaning the current training epoch is 1, then... The value is 0;

[0065] The final loss is the sum of the weighted multi-level supervision loss and the main output loss:

[0066]

[0067] in, λ is the main output loss corresponding to the last stage of the decoder, and λ is a balancing hyperparameter that determines the relative importance of the main output loss.

[0068] The training process uses the AdamW optimizer with an initial learning rate of 0.001; at the same time, a cosine annealing learning rate scheduling strategy, L2 regularization, and early stopping strategy are adopted.

[0069] A 2D medical image segmentation system based on Mamba and UNet includes a memory and one or more processors. The memory stores executable code, and when the one or more processors execute the executable code, they implement the aforementioned 2D medical image segmentation method.

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

[0071] 1. Superior Fine-Scale Segmentation Capability: Through the innovative Local-Global Fusion Selective State Space 2D (LGF-SS2D) module combined with multi-scale convolution operations, the model can simultaneously and efficiently capture global contextual information and enhance local detail features. The Dynamic Gated Fusion Block (QuadGate) further realizes the adaptive fusion of state space features and multi-scale local features, significantly improving the model's segmentation accuracy for fine structures such as small lesions and complex organ boundaries.

[0072] 2. Optimized Cross-Scale Feature Fusion: The proposed Multi-Level Cross-Scale Feature Fusion Block (MCFB) effectively integrates feature maps from all encoder stages into the current decoder layer through cross-layer attention gating and multi-scale feature transformation. This design not only fuses low-level spatial details and high-level semantic information, but also strengthens the feedback guidance from high-level semantics to low-level features, addressing the shortcomings of existing methods in cross-scale feature fusion and thus obtaining more accurate segmentation boundaries.

[0073] 3. Adaptive Optimization Training Strategy: An adaptive hierarchical loss function based on gradient statistics is introduced. By analyzing the gradient magnitude of the output at each decoder stage in real time, the weights of the supervised loss at each level are dynamically adjusted. This adaptive weight allocation mechanism can respond to changes in the learning needs of different training stages and decoder levels, overcoming the limitations of fixed weights or simple learnable weight strategies, and effectively improving the model's convergence speed and final performance.

[0074] 4. Efficient long-range dependency modeling: The core state-space model component (Mamba) provides long-range dependency modeling capabilities with linear computational complexity, overcoming the limitations of traditional CNNs in long-range modeling, while avoiding the burden of secondary computational complexity of Transformer.

[0075] 5. Excellent overall segmentation performance: Experimental validation on publicly available medical image datasets such as Synapse, ACDC, ISIC2017, ISIC2018, and CVC-ClinicDB shows that the LGFVM-UNet model proposed in this invention achieves leading or highly competitive results in key evaluation metrics such as DSC, HD95, mean intersection-over-union ratio, accuracy, sensitivity, and specificity, demonstrating its efficiency, accuracy, and robustness in handling medical image segmentation tasks of different modalities and anatomical locations.

[0076] 6. Highly practical and with broad application prospects: The model has a clear structure and an effective training strategy, enabling it to automatically and intelligently complete various 2D medical image segmentation tasks, such as organ and lesion region segmentation. It provides a high-precision auxiliary tool for clinical diagnosis, treatment planning, and surgical navigation, and has significant clinical value and broad application potential. Attached Figure Description

[0077] Figure 1 This is a diagram of the overall architecture of the present invention.

[0078] Figure 2 This is a structural diagram of LGF-VSS and QuadGate in this invention.

[0079] Figure 3 This is a structural diagram of LGF-SS2D in this invention.

[0080] Figure 4 This is a structural diagram of the MCFB in this invention. Detailed Implementation

[0081] The present invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be noted that the embodiments described below are intended to facilitate the understanding of the present invention and do not constitute any limitation thereof.

[0082] like Figure 1 As shown, a 2D medical image segmentation method based on Mamba and UNet includes the following steps:

[0083] Step 1: Medical image segmentation dataset collection and preprocessing.

[0084] The method proposed in this invention has been validated on five publicly available medical image datasets, including:

[0085] 1. Synapse multi-organ segmentation dataset;

[0086] 2. ACDC cardiac MRI dataset;

[0087] 3. ISIC2017 Skin Lesion Dataset;

[0088] 4. ISIC2018 Skin Lesion Dataset;

[0089] 5. CVC-ClinicDB colonoscopy polyp dataset.

[0090] The aforementioned dataset spans multiple modalities, including different organ sites in medical images. Validation on the aforementioned dataset demonstrates that the method proposed in this invention is effective for any two-dimensional medical image dataset.

[0091] For medical image datasets, the method proposed in this invention requires a unified preprocessing operation, the implementation process of which is as follows:

[0092] 1. Image size normalization: The medical images in the dataset are adjusted to the same resolution and then normalized, adjusting each pixel to the range [0,1]. In the five previously mentioned disclosed medical image dataset embodiments of this invention, the resolution of each medical image is adjusted to 256×256.

[0093] 2. Image Data Augmentation: Perform data augmentation on each image in the dataset, including random clockwise rotation within a degree of 75° to 105°, elastic deformation with a 30% probability, and applying a random change in the brightness value of each pixel within ±20% of the original brightness value.

[0094] 3. Training and test set partitioning: The data augmented dataset was divided into training and test sets in a 4:1 ratio. The training set was used to train the network model, enabling LGFVM-UNet to learn features autonomously and select and adjust parameters. The test set was used to evaluate the final generalization error.

[0095] Step 2: Construct a 2D medical image segmentation model based on Mamba and UNet, namely LGFVM-UNet. The specific implementation scheme of this segmentation network is as follows:

[0096] 1. Implementation of the overall network architecture:

[0097] The network employs a U-shaped encoder-decoder architecture, comprising a four-stage encoder and a four-stage decoder. For the input image... The network is processed through four main stages: initial feature embedding, multi-scale encoding via an encoder composed of LGF-VSS, progressive decoding via a decoder composed of LGF-VSS and MCFB, and prediction generation.

[0098] 2. Implementation of initial feature embedding:

[0099] For the input image This invention uses a patch embedding layer for initial feature embedding, transforming the original pixel space into a low-dimensional semantic space, significantly reducing computational complexity. Simultaneously, it employs convolutional kernels to extract local features, preserving the spatial structure information of the image and providing an effective input representation for the encoder. This process can be represented as:

[0100]

[0101] Among them, Conv 4×4 This represents a convolution with a kernel size of 4, 3 input channels, and 96 output channels. After the above operations, for By performing a dimension permutation, the height and width dimensions are moved forward to before the channel dimensions, resulting in a feature map that can be input into the first stage of the encoder. Compared to the input image The channel dimension and the order of the height and width dimensions are different, the batch processing dimension remains the same, the number of channels becomes 96, and the resolution size becomes 1 / 4 of the original.

[0102] 3. Encoder Implementation:

[0103] The encoder comprises four stages, each extracting hierarchical features through two Local-Global Fusion Visual State Space Blocks (LGF-VSS). Except for the final stage, all stages of the encoder also include a PatchMerging layer. The LGF-VSS also includes a Dynamically Gated Fusion Block (QuadGate) and a Local-Global Fusion Selective State Space 2D Module (LGF-SS2D). Their specific implementation steps are as follows:

[0104] (1) For the encoder at stage i (where i takes the value of 1, 2, 3, or 4), the final output of the encoder or block embedding layer in the previous stage is denoted as the feature map. Preprocess it:

[0105]

[0106] Among them, Conv 3Δ3This represents a convolution with a kernel size of 3 and padding of 1. LN(·) denotes layer normalization, and σ is the SiLU activation function. After the above operations, the feature map is obtained.

[0107] (2) For feature maps The four sets of spatially adaptive fusion weights are obtained through QuadGate, such as... Figure 2 The relevant sections of QuadGate are shown below. The various operations performed in QuadGate can be represented by the following formulas:

[0108]

[0109] Among them, AvgPool is an average pooling layer that can pool the input feature map. The spatial resolution is reduced to 1×1. Linear layers are linear layers with the same number of input and output channels, used for cross-channel interaction of feature maps without changing the number of channels in the image. Conv 1Δ1 This represents a 1×1 convolution, which expands the channel dimension by a factor of 4. The resulting gated tensor Divided along the channel dimension into four components {ω1, ω3, ω5, ω m Each component has a shape of B×C×1×1, which is used for fusion in the LGF-SS2D component. Softmax represents the Softmax function, which is used to normalize the weights in the components to ensure that the sum of the weight values ​​on each channel is 1.

[0110] (3) For feature maps This is achieved by fusing multi-scale local features with the global state space using LGF-SS2D. For example... Figure 3 As shown, the LGF-SS2D process can be represented as:

[0111]

[0112] Here, ⊙ represents broadcast element-wise multiplication, which means that the component of shape B×C×1×1 is expanded to shape B×C×H×W through a broadcast mechanism and then multiplied element-wise with the feature map of shape B×C×H×W. SS2D is the state space module, and Conv... k×k For a convolution kernel of size k, padding is... The convolution is BN for batch normalization.

[0113] The SS2D process can be represented as:

[0114]

[0115] F = Reshape(Transpose(Y))

[0116] Here, Mamba represents the Mamba state-space model, and Transpose represents dimensional permutation. Replace with Reshape represents tensor reshaping, which will... Remodeling

[0117] (4) Optimize the gradient flow using double residual connections. This process can be expressed as:

[0118]

[0119] Where LN(·) denotes layer normalization, σ is the SiLU activation function, and ⊙ denotes broadcast element-wise multiplication. The final output feature map is This concludes the implementation of LGF-VSS within a single encoder stage. Since LGF-VSS operations do not alter the size of the input feature map, multiple LGF-VSS operations can be stacked for multi-step feature extraction. In this invention, all encoder stages contain only two LGF-VSS operations.

[0120] For the final encoder stage, this marks the end of all encoder stages. That is, the final output feature map of the encoder at the current stage can be used as the input of the decoder in the first stage and subsequent MCFB, while for the encoder in other stages, Patch Merging is still required to obtain the final output feature map, which serves as the input for the next stage encoder and subsequent MCFB.

[0121] (5) Patch Merging, which downsamples the feature maps of the encoder except for the last stage encoder through spatial reconstruction and channel compression, thereby reducing resolution while enhancing feature representation. The process can be represented as:

[0122]

[0123] Where LN(·) represents layer normalization, Conv 1×1 This indicates that the kernel size is 1 and the number of input channels is 4. The number of channels is twice the number of output channels. The convolution of the number of channels, with PixleUnshuffle(·) as the pixel recombination operation, reassembles the pixels of each 2×2 local region in the input tensor into the channel dimension, achieving lossless downsampling that halves the spatial resolution while increasing the number of channels by four times. After the above operations, one encoder stage is completed, and the output feature map of the current encoder stage is finally obtained.

[0124] 4. Decoder Implementation:

[0125] The decoder consists of four stages. All stages except the final one include two Local-Global Fusion Visual State Space Blocks (LGF-VSS) for feature refinement, while the final encoder stage contains only one LGF-VSS. Each stage's encoder feeds the feature map refined by the LGF-VSS into a Multi-Scale Feature Fusion Block (MCFB), such as... Figure 4 As shown, MCFB first upsamples it, and then fuses it with the feature maps from all stages of the encoder output to generate a feature map that contains both high-level semantic understanding and precise spatial positioning information.

[0126] The LGF-VSS implementation in the decoder is the same as in the encoder, while the specific implementation steps of MCFB are as follows (it should be noted that B, C, H, W, and their changes representing the feature map shape in each step below are only valid in the current step and are used to illustrate the change in the feature map shape in each step; they cannot be used to compare the feature map shape size across steps):

[0127] (1) Feature scale alignment:

[0128] It should be noted that for the i-th decoder stage, MCFB receives feature maps from all encoder stages. and the output of the feature map in the current decoder after passing through the LGF-VSS module. And first, it is upsampled:

[0129]

[0130] Where Linear represents a linear layer with twice the number of output channels, and Reshape represents tensor reshaping. Remodeling LN(·) represents layer normalization. This yields the upsampled feature map.

[0131] To explicitly represent the differentiated processing of features at different scales, the output of the j-th stage encoder is represented as... For feature map outputs with different resolutions from encoders at different stages, MCFB first performs feature scale alignment, which can be represented as:

[0132]

[0133] Where D i-j (·) represents the feature map Cross-layer downsampling, using a convolution kernel size of 3 and an output channel number of Number of channels, step size 2 i-j The convolution with padding of 1 is complete, U j-i (·) represents the feature map Bilinear upsampling with a scaling factor of 2 j-i After the above operations, the feature map scales of different scales were successfully aligned to the upsampled feature map. Scale, i.e., the scale of the aligned feature map. After feature scale alignment, the outputs of the four different stages of the encoder will actually result in four encoders with the same scale but different values. Therefore, the following text is about The operation is obtained by aligning the feature scale of the encoder at any stage. All of these must be implemented.

[0134] (2) Spatial attention calculation:

[0135] Spatial attention module computes decoder features D i and aligned encoder features The correlation between them can be represented as follows:

[0136]

[0137] Where BN(·) represents batch normalization, Conv 1×1 This represents a convolution with a kernel size of 1. 3×3 This represents a convolution with a kernel size of 3, and Simoid represents the Sigmoid activation function. This process yields the spatial attention tensor.

[0138] (3) Channel attention calculation:

[0139] The channel attention component further captures semantic dependencies through a squeeze incentive mechanism:

[0140] Ψ channel =Sigmoid(Conv 1×1 (δ(Conv 1×1 (AvgPool(Ψ spatial )))))

[0141] Among them, AvgPool is an average pooling layer that can pool the input Ψ spatial The spatial resolution is reduced to 1×1, Conv 1×1 This is a convolution with a kernel size of 1, where δ represents the ReLU activation function and sigmoid represents the sigmoid activation function. This process yields the channel attention tensor.

[0142] (4) Comprehensive attention map generation and feature refinement:

[0143] The spatial and channel attention weights are combined to generate a comprehensive attention map:

[0144]

[0145] Where ⊙ represents broadcast element-wise multiplication, i.e., the spatial attention tensor Ψ of shape B×1×1×C. spatial and the integrated attention tensor of shape B×H×W×1 After being expanded to a shape of B×H×W×C via a broadcast mechanism, it is aligned with the channel attention tensor of shape B×H×W×C and the scale of the aligned feature map. Perform element-wise multiplication, Conv 1×1 This is a convolution with a kernel size of 1 and an output channel of 1. Batch normalization (BN) is used, and sigmoid represents the sigmoid activation function. The final result is a refined feature set.

[0146] (5) Multi-scale feature aggregation:

[0147] Refined features from the encoder stage are obtained through a spatial-channel dual attention mechanism. Then, MCFB integrates them with the current decoder features. All cross-scale features are concatenated and processed through LGF-VSS blocks, and finally combined with the original decoder features via residual connections. This process can be represented as:

[0148]

[0149] in These refer to the different refined features obtained after the outputs of encoders from lower stages, the same stage, and higher stages have undergone the aforementioned feature scale alignment, spatial attention calculation, channel attention calculation, and comprehensive attention map generation and feature refinement. Concat(·) represents concatenating all refined features along the channel dimension. Represents the features after splicing The signal is fed into the LGF-VSS mentioned above. After the above operations, the final output of the i-th stage decoder is obtained.

[0150] 5. Implementation of the final prediction generation layer:

[0151] The final prediction generation layer is used to generate feature maps from the decoder output in the final stage. Final spatial resolution upsampling and pixel-level classification are performed to obtain pixel-level mask segmentation results. The specific implementation steps are as follows:

[0152] (1) Apply a linear transformation to the low-resolution, high-channel-count feature map output by the decoder to expand the number of channels of the input features to 16 times the original number of channels, thus preparing sufficient information for subsequent spatial upsampling. This process can be represented as:

[0153]

[0154] Linear is a linear layer with four times the number of output channels as input channels.

[0155] (2) Perform layer normalization and tensor reshaping operations on the feature map after channel expansion. This process can be represented as:

[0156]

[0157] Where Reshape represents tensor reshaping, which... Remodeling

[0158] LN(·) represents layer normalization.

[0159] (3) Input the normalized high-resolution feature map into a 1×1 convolutional layer, and apply the Softmax function to map the number of channels in the feature map to the number of classes in the target segmentation task. This process can be represented as:

[0160]

[0161] Where Conv 1×1 This indicates a convolution with a kernel of 1 and an output channel count equal to the number of target segmentation categories. Softmax represents the Softmax function, which converts the value at each pixel location into a probability distribution, indicating the probability that the pixel belongs to each category. Finally, the channel with the highest probability at each pixel location in Mask' is taken as the final predicted category index for that pixel, thus generating a complete pixel-level semantic segmentation mask image, Mask.

[0162] Step 3: Train the LGFVM-UNet model on the training set using an adaptive hierarchical loss function based on gradient statistics and certain training settings, and fine-tune the network model parameters to achieve the optimal network performance.

[0163] The gradient statistics adaptive hierarchical loss balancing strategy proposed in this invention adopts different weights for the loss of different decoder layers, adapting to the constantly changing learning dynamics throughout the training process, thereby achieving better learning results.

[0164] Specifically, the gradient statistics adaptive hierarchical loss balancing strategy proposed in this invention automatically adjusts the contribution of each supervision layer based on real-time gradient statistics, that is, for the final output of each decoder stage i. With supervision, these outputs need to be processed by an auxiliary segmentation head (i.e., a convolution with a kernel size of 1 and an output channel number equal to the number of target segmentation task categories) and a Softmax function to generate a complete pixel-level semantic segmentation mask. The specific implementation steps are the same as those in step 2, which is the implementation steps for generating a complete pixel-level semantic segmentation mask in the final prediction generation layer.

[0165] The specific implementation steps of the gradient statistics adaptive hierarchical loss balancing strategy are as follows:

[0166] (1) Gradient magnitude calculation:

[0167] For the l-th decoder stage with training epoch t, ​​t≥1, first calculate the gradient magnitude:

[0168]

[0169] Where |θ| represents the total number of network parameters, This is the loss for the l-th decoder stage. The magnitude of this gradient reflects the current learning state and optimization potential of each supervised layer.

[0170] (2) Momentum smoothing weight update:

[0171] To ensure stable weight updates and prevent drastic oscillations, a momentum-based smoothing method is used to calculate the weights for the current decoder stage:

[0172]

[0173] Where α∈[0,1] is the momentum coefficient controlling the retention of historical information, exp represents the exponential function, and τ>0 is used as a temperature parameter to control the sharpness of the weight distribution. If it does not exist, that is, if the current training epoch is 1, then... The value is 0.

[0174] (3) Final composite loss:

[0175] The final composite loss in any epoch is calculated by combining the adaptively weighted supervision loss with the main output loss:

[0176]

[0177] Where λ is a balancing hyperparameter that determines the relative importance of the main output loss. and The calculation formula is the same, but the implementation steps are slightly different, because The decoder with lower resolution outputs a complete pixel-level semantic segmentation mask. Therefore, in subsequent calculations, the real segmentation mask needs to be downsampled to the same resolution using bilinear methods before processing. This step is not required in the implementation process. represent and but The following formula can be used for calculation:

[0178]

[0179] Where α and β are the cross-entropy loss values. Dice similarity coefficient loss The relative importance balancing hyperparameters, where N is the total number of image pixels, C is the number of object segmentation task categories, and y i,c This is the one-hot encoding of the pixel on the real segmentation mask or the real segmentation mask after downsampling to the same resolution. It is 1 only if the pixel belongs to class c, and 0 otherwise. This is the one-hot encoding of the pixel on the model prediction mask map, where σ is the Softmax function and ε is a very small smoothing coefficient with a value of 0.000001 to avoid division by zero.

[0180] Based on the gradient statistical adaptive hierarchical loss balancing strategy proposed in this invention, the training method is implemented in Python using the PyTorch deep learning framework. Experiments were conducted on a high-performance server equipped with an NVIDIA GeForce RTX4090 GPU, running Ubuntu 20.04. Training was completed after a certain time and several epochs.

[0181] The specific training strategy involves training the LGFVM-UNet model using the training set and employing the AdamW optimizer. To avoid overfitting, learning rate scheduling, early stopping, and L2 regularization are used. The specific implementation is as follows:

[0182] 1. Learning rate scheduling: During the training process, a cosine annealing learning rate scheduling strategy is adopted. The initial learning rate is set to 0.001. Over several epochs of the total training cycle, the learning rate gradually decreases from the initial value to the minimum learning rate set in Table 1 according to the cosine function.

[0183] 2. Early stopping strategy: If there is no improvement in performance on the validation set within 20 epochs, training is stopped early to avoid overfitting.

[0184] 3. L2 Regularization: L2 regularization (weight decay) is used to prevent overfitting. The regularization parameters are gradually adjusted during model training to ensure optimal regularization results.

[0185] Step 4: Use the test set to test and evaluate the performance of the obtained network model.

[0186] This embodiment trains the LGFVM-UNet model on test sets divided from the five publicly available medical image datasets mentioned in step 1, and then evaluates the performance of the LGFVM-UNet model in medical image segmentation using different metrics. All metrics involved in the evaluation in this embodiment are as follows:

[0187] 1. Dice similarity coefficient (DSC): Used to measure the volume overlap between the predicted and the true label.

[0188] 2.95% Hausdorf distance (HD95): Used to evaluate the accuracy of boundary delineation.

[0189] 3. Mean Intersection over Union (mIoU), Accuracy (Acc), Sensitivity (Sen), and Specificity (Spe): These metrics are used to comprehensively evaluate the overall performance of image segmentation, classification accuracy, and the model's ability to identify target structures (positive class) and background (negative class).

[0190] The results of different experiments conducted on the five publicly available medical image datasets mentioned in step 1 are shown in Tables 1, 2, 3, 4, and 5.

[0191] Table 1

[0192] Table 2

[0193]

[0194] Table 3

[0195]

[0196] Table 4

[0197]

[0198] Table 5

[0199]

[0200] The model demonstrated excellent and stable performance in evaluations on the ISIC2017, ISIC2018, CVC-ClinicDB, Synapse, and ACDC datasets. Its comprehensive metrics (such as DSC generally exceeding 91%, and class average DSC reaching 88.74% and 92.40%) and other key metrics (such as mIOU, ACC, Sen, and Spe) remained at high levels, fully validating its robustness and generalization ability. Specifically, the LGFVM-UNet model achieves a class-average DSC of 88.74% and an HD95 of 6.65 on the Synapse dataset; a class-average DSC of 92.40% and an HD95 of 21.97 on the ACDC dataset; a DSC of 91.70% and an mIoU of 84.67% on the ISIC2017 dataset; a DSC of 91.35% and an mIoU of 84.09% on the ISIC2018 dataset; and a DSC of 91.03% and an mIoU of 83.55% on the CVC-ClinicDB dataset. This demonstrates that the LGFVM-UNet model can effectively segment medical images across different modalities and regions, exhibiting high diagnostic accuracy and reliability.

[0201] Through the above tests and evaluations, the LGFVM-UNet model proposed in this invention demonstrates excellent segmentation performance in medical images, enabling automatic and intelligent medical image segmentation. The model's efficiency and accuracy provide a reliable auxiliary diagnostic tool for clinical practice, exhibiting broad application prospects and significant clinical value.

[0202] The embodiments described above provide a detailed explanation of the technical solutions and beneficial effects of the present invention. It should be understood that the above descriptions are merely specific embodiments of the present invention and are not intended to limit the present invention. Any modifications, additions, and equivalent substitutions made within the scope of the principles of the present invention should be included within the protection scope of the present invention.

Claims

1. A 2D medical image segmentation method based on Mamba and UNet, characterized in that, Includes the following steps: (1) Collect medical image segmentation datasets and preprocess them to obtain training sets; (2) Construct a 2D medical image segmentation model based on Mamba and UNet. The 2D medical image segmentation model includes a block embedding layer, an encoder, a decoder and a prediction generation layer. The block embedding layer is used to map the input image into a semantic feature vector, which is then input into the encoder; The encoder consists of multiple stages. The first stage receives the output of the block embedding layer through the LGF-VSS module to extract feature maps. Each stage other than the first stage receives the output of the previous stage through the LGF-VSS module to extract feature maps. In each stage except the final stage, the feature maps extracted by the LGF-VSS module are downsampled through the corresponding block merging layer before being input into the next stage; In the final stage, the feature maps extracted by the LGF-VSS module are directly input into the decoder without passing through the block merging layer; The decoder consists of multiple stages. The first stage uses the LGF-VSS module to receive the output of the last stage of the encoder for feature refinement. Each stage except the first stage uses the LGF-VSS module to receive the output of the previous stage for feature refinement. Each stage uses the MCFB module to fuse the refined feature map with the feature maps from all stages of the encoder to generate a feature map that contains both high-level semantic understanding and precise spatial positioning information. The feature maps generated in each stage except the last stage are input into the next stage, and the feature map generated in the last stage is input into the prediction generation layer. The prediction generation layer upsamples and performs pixel-level classification on the feature map output from the last stage of the decoder to obtain pixel-level mask segmentation results. The LGF-VSS module includes the LGF-SS2D module and a dynamic gating fusion block; The LGF-SS2D module captures the global context through four-directional state space scanning to obtain state space features, and extracts multi-scale local features through multi-scale convolutions with kernel sizes of 1, 3, and 5. The dynamic gated fusion block generates four sets of spatially adaptive fusion weights, which are weighted and fused with the state space features and multi-scale local features output by the LGF-SS2D module to obtain fused features. The fused features are combined with the original input of the LGF-SS2D module through a dual residual connection mechanism, and the feature map is output after layer normalization. The MCFB module fuses the refined feature map with feature maps from all stages of the encoder. Specifically, it unifies the resolution of feature maps from all stages of the encoder by upsampling or downsampling. Generate spatial-channel dual attention weights and weight features from all stages of the encoder; The weighted encoder features are concatenated with the decoder features of the current stage, fused by LGF-VSS, and then output through residual connection. (3) Design an adaptive hierarchical loss function based on gradient statistics to train the 2D medical image segmentation model on the training set; (4) Input the segmented medical image into the trained model to complete the image segmentation.

2. The 2D medical image segmentation method based on Mamba and UNet according to claim 1, characterized in that, In step (2), both the encoder and decoder contain four stages.

3. The 2D medical image segmentation method based on Mamba and UNet according to claim 1, characterized in that, The LGF-SS2D module captures the global context through a four-directional state space scan, specifically as follows: The feature map is unfolded into a sequence along four dimensions: forward height, backward height, forward width, and backward width. The Mamba model was used to model long-range dependencies of sequences in each direction. By fusing the output features from the four directions, we obtain the global context features, i.e., the state space features.

4. The 2D medical image segmentation method based on Mamba and UNet according to claim 1, characterized in that, The specific process of generating four sets of spatial adaptive fusion weights by the dynamic gated fusion block is as follows: Global average pooling is performed on the features input to the LGF-VSS module to obtain channel statistics; Channel statistics are used to generate four sets of channel attention weights through a linear layer, a SiLU activation function, and a 1×1 convolution. The weights are normalized using the Softmax function to obtain four sets of spatially adaptive fusion weights.

5. The 2D medical image segmentation method based on Mamba and UNet according to claim 1, characterized in that, In step (2), the block merging layer works as follows: The input feature map is divided into 2×2 local regions, and pixels are reorganized to the channel dimension to achieve a 4-fold channel expansion. The number of channels is compressed to twice the size using 1×1 convolution; The application layer normalizes the output downsampled feature map.

6. The 2D medical image segmentation method based on Mamba and UNet according to claim 1, characterized in that, In step (3), an adaptive hierarchical loss function based on gradient statistics is designed, specifically including: Calculate the gradient magnitude of the decoder output at each stage except the last stage: in, This represents the total number of network parameters. It is the decoder except for the last stage. Losses at each stage; The loss weights for each stage are updated using momentum smoothing: in, It is the momentum coefficient that controls the retention of historical information, and exp represents the exponential function. The sharpness of the weight distribution is controlled by the temperature parameter; if If it does not exist, meaning the current training epoch is 1, then... The value is 0; The final loss is the sum of the weighted multi-level supervision loss and the main output loss: in, It is the main output loss corresponding to the last stage of the decoder. It is a balancing hyperparameter that determines the relative importance of the main output loss.

7. The 2D medical image segmentation method based on Mamba and UNet according to claim 1, characterized in that, In step (3), the AdamW optimizer is used during the training process, with an initial learning rate of 0.

001. At the same time, a cosine annealing learning rate scheduling strategy, L2 regularization, and early stopping strategy are adopted.

8. A 2D medical image segmentation system based on Mamba and UNet, characterized in that, The device includes a memory and one or more processors, wherein the memory stores executable code, and the one or more processors execute the executable code to implement the 2D medical image segmentation method according to any one of claims 1-7.

Images