2D medical image segmentation method and system based on Mama and UNet

Abstract

The invention discloses a 2D medical image segmentation method and system based on Mama and UNet, and the method comprises the steps: collecting and preprocessing a medical image segmentation data set, and obtaining a training set; constructing a 2D medical image segmentation model based on Mama and UNet, wherein the 2D medical image segmentation model comprises 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 with segmentation into the trained model to complete image segmentation. According to the invention, the method can achieve the automatic and intelligent segmentation of the medical image through the innovative construction of the 2D medical image segmentation model based on Mamba and UNet, and is higher in segmentation accuracy and efficiency.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of computer vision deep learning and medical image processing, and particularly relates to a 2D medical image segmentation method and system based on Mamba and UNet. BACKGROUND

[0002] Medical image segmentation is a key task in computer vision, which can accurately divide anatomical structures and pathological regions, thereby improving diagnosis, treatment planning and surgical navigation. Unlike conventional images, medical images have unique challenges such as low contrast, fuzzy boundaries, inter-class similarity and noise artifacts, making accurate segmentation particularly difficult.

[0003] In the prior art, convolutional neural networks (CNN), especially UNet, have become the basic architecture for medical image segmentation. However, the inherent limitations of CNN in modeling long-range dependencies limit its performance on complex anatomical structures. The Transformer largely solves this limitation through the self-attention mechanism, but its quadratic computational complexity and difficulty in preserving fine local features required for accurate boundary delineation are still problematic.

[0004] With the development of state space models (SSM), especially Mamba, due to its linear computational complexity and efficient long-range modeling capability, it has become a promising alternative.

[0005] A brain tumor image segmentation method based on multi-scale convolution and Mamba structure is disclosed in Chinese patent document No. CN118447244A, which reduces the computational complexity by introducing a state space model and improves the recognition ability of detailed features by introducing a cascade residual multi-scale convolution module. Chinese patent document No. CN120070891A discloses a tumor segmentation method based on Mamba guided multi-encoder fusion, which includes: 1. constructing a multi-encoder segmentation model that combines Mamba and convolutional neural networks (CNN); 2. training and optimizing the model 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 graph of the tumor segmentation result.

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

[0007] 1. The existing Mamba-based 2D medical image segmentation method is generally not ideal for segmentation of fine structures such as small lesions or complex organ boundaries.

[0008] 2. The existing Mamba-based 2D medical image segmentation method with multiple encoders and multiple decoders often has defects in cross-scale feature fusion between encoders and decoders, and lacks feedback guidance from high-level semantics to low-level features.

[0009] 3. The existing Mamba-based 2D medical image segmentation method usually uses fixed weights or simple learnable weights if using multi-level loss combination, which cannot adapt to the dynamic training needs of different decoder levels. SUMMARY

[0010] In view of the shortcomings of the existing medical image segmentation method, the present application provides a 2D medical image segmentation method and system based on Mamba and UNet, which realizes automatic and intelligent segmentation of medical images through the innovative architecture design of local-global fusion visual Mamba UNet network (LGFVM-UNet), and has high segmentation accuracy and efficiency.

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

[0012] (1) Collecting medical image segmentation data set and pre-processing to obtain training set;

[0013] (2) Constructing a 2D medical image segmentation model based on Mamba and UNet, wherein the 2D medical image segmentation model comprises a patch embedding layer, an encoder, a decoder and a prediction generation layer;

[0014] The patch embedding layer is used to map the input image to a semantic feature vector and then input to the encoder; the specific operation is to input the medical image through a convolution layer with a 4x4 size convolution kernel;

[0015] The encoder comprises multiple stages, the first stage receives the feature map extracted by the output of the patch embedding layer through a local-global fusion selective state space 2D module (LGF-VSS), and each stage after the first stage receives the feature map extracted by the output of the previous stage through the LGF-VSS module; in each stage except the last stage, the feature map extracted by the LGF-VSS module is input to the next stage after being down-sampled by the corresponding block merging layer; in the last stage, the feature map extracted by the LGF-VSS module is directly input to the decoder without passing through the block merging layer;

[0016] The decoder comprises multiple stages, the first stage receives the output of the last stage of the encoder through the LGF-VSS module for feature refinement, and each stage after the first stage receives the output of the previous stage through the LGF-VSS module for feature refinement; each stage fuses the refined feature map with the feature map from the output of all stages of the encoder through a multi-level cross-scale feature fusion block (MCFB) to generate a feature map containing both high-level semantic understanding and accurate spatial positioning information; the feature map generated in each stage except the last stage is input into the next stage, and the feature map generated in the last stage is input into a prediction generation layer;

[0017] The prediction generation layer performs upsampling and pixel-level classification on the feature map output by the last stage of the decoder to obtain a pixel-level mask segmentation result.

[0018] (3) An adaptive hierarchical loss function based on gradient statistics is designed, and a 2D medical image segmentation model is trained on a training set;

[0019] (4) The medical image with segmentation is input into the trained model to complete the segmentation of the image.

[0020] In the present application, in step (2), the encoder and the decoder each comprise four stages.

[0021] In step (2), the LGF-VSS module efficiently fuses local detailed features and global context information by combining a selective state space model and a multi-scale convolution operation, and enhances the feature representation capability of the model in the visual task. Specifically, the LGF-VSS module comprises an LGF-SS2D module and a dynamic gating fusion block.

[0022] The LGF-SS2D module captures global context through four-direction state space scanning to obtain state space features, and jointly extracts multi-scale local features through multi-scale convolutions with kernel sizes of 1, 3 and 5.

[0023] The dynamic gating fusion block generates four groups of spatial adaptive fusion weights, and fuses the state space features and the multi-scale local features output by the LGF-SS2D module to obtain fusion features.

[0024] The fusion features are combined with the original input of the LGF-SS2D module through a double residual connection mechanism, and the output feature map is output through layer normalization.

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

[0026] Layer Normalization (LayerNorm) is applied to the input feature map for standardization, and the feature dimension order is adjusted from (B, H, W, C) to (B, C, H, W), then the input features are preliminarily spatially extracted using convolution, and non-linear representation is introduced through SiLU activation function to obtain the primary feature map in LGF-VSS;

[0027] The primary feature map in LGF-VSS is input into LGF-SS2D to obtain three convolution branch feature maps and one state space feature map;

[0028] The primary feature map in LGF-VSS is input into a dynamic gate fusion block (QuadGate) to obtain four groups of spatially adaptive fusion weights;

[0029] The state space feature map and the three convolution branch feature maps are weighted and summed according to the weights to realize adaptive fusion of local and global features;

[0030] The fused features and the original LGF-VSS input are combined through residual connection, the output features are normalized by LayerNorm, and a gating mechanism is used, i.e., multiplied by the original input to control information flow and enhance feature selectivity. Finally, the processed feature map dimension is restored from (B, C, H, W) to (B, H, W, C) to maintain the same dimension order as the input.

[0031] The LGF-SS2D module captures the global context through four-direction state space scanning, the process is as follows:

[0032] The feature map is unfolded into a sequence along the height forward, height reverse, width forward, and width reverse dimensions;

[0033] The Mamba model is used to model long-range dependencies for each direction sequence;

[0034] The output features of the four directions are fused to obtain the global context feature, i.e., the state space feature.

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

[0036] The SS2D module realizes four-direction selective state space scanning of the primary feature map in LGF-VSS, i.e., the feature map is unfolded into a sequence along the height forward, height reverse, width forward, and width reverse dimensions. Then the Mamba model is used to model long-range dependencies for each direction sequence, and finally the output features of the four directions are fused to capture the global spatial context information of the feature map and obtain the state space feature map;

[0037] While the SS2D module and the Mamba model capture the global spatial context information of the feature map, the 1x1 convolution branch, the 3x3 convolution branch and the 5x5 convolution branch are used in parallel to extract local features of different receptive fields from the primary feature map in the LGF-VSS. After each branch, the output of the branch is input into the LayerNorm and SiLU activation to enhance the feature expression capability and obtain three convolution branch feature maps.

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

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

[0040] The channel statistical information is used to generate four sets of channel attention weights through a linear layer, a SiLU activation function and a 1x1 convolution;

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

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

[0043] The input feature map is divided into 2x2 local regions, and the pixels are reorganized into the channel dimension to achieve 4 times channel expansion;

[0044] The channel number is compressed to 2 times through a 1x1 convolution;

[0045] The output down-sampled feature map is obtained by applying the LayerNorm.

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

[0047] The resolution of the feature maps output from all stages of the encoder is unified through up-sampling or down-sampling;

[0048] Spatial-channel dual attention weights are generated to weight the features of all stages of the encoder;

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

[0050] Further, 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 gate mechanism is used to interact it with the encoder feature map corresponding to the current decoder layer to generate spatial and channel dual attention weights, and to weight the original feature map to highlight important information;

[0052] The spatial resolution of all encoder feature maps including the current layer is uniformly adjusted to the target resolution of the current decoder layer by upsampling or downsampling operations, while the number of channels is adjusted to one quarter of the target number of channels;

[0053] All feature maps after resolution and channel adjustment are spliced in the channel dimension to form multi-scale fusion feature maps, and deep feature map fusion and enhancement are performed through LGF-VSS;

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

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

[0056] Linear transformation is applied to the low-resolution and high-channel number feature maps output by the decoder to expand the number of input feature channels to four times the original number of channels, preparing sufficient information for subsequent spatial upsampling.

[0057] Layer normalization and tensor reshaping operations are performed on the feature maps after channel expansion. The normalized high-resolution feature maps are input into a 1x1 convolutional layer to map the number of feature channels to the number of classes of the target segmentation task. Then, the Softmax function is applied to the class score vector of each pixel position output by the 1x1 convolution to convert the value at each pixel position to 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 of the pixel, thereby generating a complete pixel-level semantic segmentation mask map.

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

[0059] The gradient amplitude of the output of each stage of the decoder except the last stage is calculated:

[0060]

[0061] where |θ| represents the total number of network parameters, is the loss of the lth stage of the decoder except the last stage;

[0062] The momentum smoothing is used to update the loss weights of each stage:

[0063]

[0064] where α∈[0,1] is the momentum coefficient controlling the retention of historical information, exp represents the exponential function, τ>0 is the temperature parameter controlling the sharpness of the weight distribution; if If not, i.e., the current training epoch is 1, then is 0;

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

[0066]

[0067] wherein, is the main output loss corresponding to the last stage of the decoder, and λ is a balance 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, the cosine annealing learning rate scheduling strategy, L2 regularization, and early stopping strategy are used.

[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 the one or more processors execute the executable code to implement the above-mentioned 2D medical image segmentation method.

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

[0071] 1. Excellent fine structure segmentation capability: Through the innovative local-global fusion selective state space 2D module (LGF-SS2D) combined with multi-scale convolution operation, the model can efficiently capture global context information and enhance local detail features. The dynamic gate fusion block (QuadGate) further realizes the adaptive fusion of state space features and multi-scale local features, significantly improving the segmentation accuracy of the model for small lesions, complex organ boundaries and other fine structures.

[0072] 2. Optimized cross-scale feature fusion: The proposed multi-level cross-scale feature fusion block (MCFB) effectively integrates all encoder stage feature maps into the current decoder layer through cross-layer attention gating and multi-scale feature transformation. This design not only integrates low-level spatial details and high-level semantic information, but also particularly strengthens the feedback guidance from high-level semantics to low-level features, solving the defects of existing methods in cross-scale feature fusion, thereby obtaining more accurate segmentation boundaries.

[0073] 3. Self-adaptive optimization training strategy: An adaptive hierarchical loss function based on gradient statistics is introduced, which dynamically adjusts the weights of each level of supervision loss by real-time analysis of the gradient amplitude of each decoder stage output. This adaptive weight allocation mechanism can respond to changes in learning needs at different training stages and decoder levels, overcoming the limitations of fixed weights or simple learnable weight strategies, effectively improving the convergence speed and final performance of the model.

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

[0075] 5. Excellent comprehensive segmentation performance: Experiments on Synapse, ACDC, ISIC2017, ISIC2018, CVC-ClinicDB public medical image datasets show that the LGFVM-UNet model proposed in the present application has achieved leading or highly competitive results in key evaluation indicators such as DSC, HD95, average intersection over union, accuracy, sensitivity and specificity, proving its efficiency, accuracy and robustness in processing different modalities, different anatomical sites medical image segmentation tasks.

[0076] 6. Strong practicability and broad application prospect: The model structure is clear, the training strategy is effective, and it can automatically and intelligently complete various 2D medical image segmentation tasks such as organ and lesion region segmentation, providing a high-precision auxiliary tool for clinical diagnosis, treatment planning and surgical navigation, and has significant clinical value and broad application potential. BRIEF DESCRIPTION OF DRAWINGS

[0077] Figure 1 is the overall architecture diagram of the present application.

[0078] Figure 2 is the structure diagram of LGF-VSS and QuadGate in the present application.

[0079] Figure 3 is the structure diagram of LGF-SS2D in the present application.

[0080] Figure 4 is the structure diagram of MCFB in the present application. DETAILED DESCRIPTION

[0081] The present application will be further described in detail below in conjunction with the drawings and examples, it should be pointed out that the following examples are intended to facilitate the understanding of the present application, and do not have any limiting effect on it.

[0082] As shown in Figure 1 , 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 the present application has been verified on five public 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 above datasets span multiple modalities, including different organ sites in medical images, and the validation of the above datasets proves that the method proposed in the present application is effective for any two-dimensional medical image dataset.

[0091] For medical image datasets, the method proposed in the present application needs to be uniformly preprocessed, and the implementation process is as follows:

[0092] 1. Image size standardization: adjust the medical images in the dataset to the same resolution, and then normalize them to adjust each pixel to the range of [0, 1]. In the present application, the resolution of each medical image is adjusted to 256x256 for the five publicly available medical image datasets mentioned above.

[0093] 2. Image data augmentation: data augmentation is performed on each image in the dataset, including randomly rotating the image clockwise by an angle within 75° to 105°, performing elastic deformation with a probability of 30%, and applying a random change to the brightness value of each pixel within ±20% of the original brightness value.

[0094] 3. Training set and test set division: the data set after data augmentation is divided into training set and test set according to the ratio of 4:1, wherein the training set is used for network model training, so that LGFVM-UNet learns features from it autonomously, and parameters are selected and adjusted; the test set is used to evaluate the final generalization error.

[0095] Step 2, construct a 2D medical image segmentation model based on Mamba and UNet, i.e. LGFVM-UNet. The specific implementation of the segmentation network is as follows:

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

[0097] The network adopts a U-shaped encoder-decoder structure, including four stages of encoder and four stages of decoder. For the input image The network processes through four main stages of initial feature embedding, multi-scale encoding through the encoder composed of LGF-VSS, progressive decoding and prediction generation through the decoder composed of LGF-VSS and MCFB.

[0098] 2. Implementation of initial feature embedding:

[0099] For the input image The present application performs initial feature embedding through a patch embedding layer, converts the original pixel space to a low-dimensional semantic space, and significantly reduces the computational complexity; at the same time, local feature extraction is realized through a convolution kernel, and the spatial structure information of the image is preserved to provide an effective input representation for the encoder. This process can be represented as:

[0100]

[0101] where Conv 4×4 represents a convolution with a kernel size of 4, and the number of input channels is 3 and the number of output channels is 96. After the above operation, the dimension of is transposed to move the height dimension and the width dimension in front of the channel dimension to obtain a feature map that can be input to the first stage of the encoder. Compared with the input image , the channel dimension and the order of the height dimension and the width dimension are different, the batch dimension is unchanged, the number of channels is changed to 96, and the resolution size is changed to 1 / 4 of the original.

[0102] 3. Implementation of the encoder:

[0103] The encoder includes four stages, each of which extracts hierarchical features through two local-global fusion visual state space blocks (LGF-VSS), and all stages of the encoder except the last stage include patch merging layers (Patch Merging). In LGF-VSS, dynamic gate fusion blocks (QuadGate) and local-global fusion selective state space 2D modules (LGF-SS2D) are also included. Their specific implementation steps are as follows:

[0104] (1) For the i-th stage (i takes a value of 1, 2, 3, or 4) of the encoder, let the final output of the previous stage of the encoder or the patch embedding layer be a feature map

[0105]

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

[0107] (2) For the feature map , spatially adaptive four-group fusion weights are obtained through a QuadGate, as shown in Figure 2 . The relevant part of the QuadGate is shown in the following formula:

[0108]

[0109] where AvgPool is an average pooling layer that can reduce the spatial resolution of the input feature map to 1x1, Linear is a linear layer with the same number of input and output channels, which is used for cross-channel interaction of the feature map without changing the channel number of the image, Conv 1Δ1 denotes a 1x1 convolution that expands the channel dimension by 4 times. The generated gating tensor is divided into four components {ω1, ω3, ω5, ω m} along the channel dimension, each with a shape of BxCx1x1, which is used for fusion in the LGF-SS2D component, and Softmax denotes a Softmax function that normalizes the weights in the component to ensure that the sum of the weight values on each channel is 1.

[0110] (3) For the feature map , multi-scale local features and global state spaces are fused through LGF-SS2D. As shown in Figure 3 , the LGF-SS2D process can be represented as:

[0111]

[0112] where ⊙ denotes broadcast element-wise multiplication, i.e., the component with a shape of BxCx1x1 is expanded to a shape of BxCxHxW through the broadcast mechanism, and then multiplied element-wise with the feature map with a shape of BxCxHxW, SS2D is a state space module, Conv k×k is a convolution with kernel size k and padding , and BN is batch normalization.

[0113] The process of SS2D can be represented as:

[0114]

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

[0116] where Mamba denotes the Mamba state space model, Transpose denotes dimension transposition, and Reshape denotes tensor reshaping. to Reshape denotes tensor reshaping. to

[0117] (4) Gradient flow is optimized using double residual connection, which can be represented 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 The implementation of LGF-VSS in the single encoder stage ends here. Since the LGF-VSS operation does not change the size information of the input feature map, multiple LGF-VSSs can be stacked for multi-step feature extraction. In the present application, all encoder stages only contain 2 LGF-VSSs.

[0120] For the last stage of the encoder, all encoder stages end here, i.e., the final output feature map of the current stage encoder, which can be used as the input of the first stage decoder and subsequent MCFB, and for other stages of the encoder, Patch Merging is still needed to obtain the final output feature map as the input of the next stage encoder and subsequent MCFB.

[0121] (5) Patch Merging, which realizes the feature map downsampling of the encoder except the last stage of the encoder through spatial reorganization and channel compression, reduces the resolution while enhancing the feature expression capability. The process can be represented as:

[0122]

[0123] where LN(·) denotes layer normalization, Conv 1×1 denotes a convolution kernel size of 1, an input channel number of 4 times the channel number, and an output channel number of 2 times the channel number, PixleUnshuffle(·) is a pixel reorganization operation, which realizes lossless downsampling by reorganizing the pixels in each 2x2 local area in the input tensor to the channel dimension, reducing the spatial resolution by half while the channel number becomes 4 times. After the above operation, one encoder stage is completed, and the output feature map of the current encoder stage is obtained

[0124] 4. Implementation of the decoder:

[0125] The decoder contains 4 stages, all of which contain 2 local-global fusion visual state space (LGF-VSS) blocks for feature refinement, except for the last stage which contains only 1 LGF-VSS block. The encoder of each stage sends the feature map refined by the LGF-VSS to a multi-level cross-scale feature fusion block (MCFB), as shown in Figure 4 The MCFB first up-samples the feature map, and then fuses it with the feature maps output from all the encoder stages to generate a feature map that contains both high-level semantic understanding and precise spatial positioning information.

[0126] The implementation of the LGF-VSS in the decoder is the same as in the encoder, and the specific implementation steps of the MCFB are as follows (it should be noted that B, C, H, W and their changes in the following steps represent the shape of the feature map, which is only valid for the current step and is used to illustrate the change of the feature map shape in each step, and 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, the MCFB receives feature maps from all encoder stages and the output of the feature map in the current decoder after the LGF-VSS module and first up-samples them:

[0129]

[0130] where Linear is a linear layer with 2 times the number of input channels as the number of output channels, Reshape represents tensor reshaping, which reshapes to LN(·) represents layer normalization. Thus, the up-sampled feature map

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

[0132]

[0133] where D i-j (·) is a cross-layer down-sampling of the feature map using a convolution kernel size of 3 and an output channel number of Channel number, step size is 2 i-j , the convolution is completed, U j-i (·) is the bilinear up-sampling of the feature map , and the scaling factor is 2 j-i . After the above operation, the feature map scale of the feature of different scales is successfully aligned to the feature map scale of the up-sampled feature , that is, the feature map scale after alignment The output of the four different stage encoders after feature scale alignment actually obtains four features with the same scale but different values Thus, the operation of below is implemented for obtained after feature scale alignment of any stage encoder.

[0134] (2) Spatial attention calculation:

[0135] The spatial attention module calculates the correlation between the decoder feature D i and the aligned encoder feature , and the process can be represented as:

[0136]

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

[0138] (3) Channel attention calculation:

[0139] The channel attention component further captures semantic dependency relationships through the squeeze-and-excitation mechanism:

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

[0141] Wherein, AvgPool is an average pooling layer, which can reduce the spatial resolution of input Ψ spatial to 1x1, Conv 1×1 is a convolution with a convolution kernel size of 1, δ represents a ReLU activation function, and Sigmoid represents a Sigmoid activation function. This process obtains a channel attention tensor

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

[0143] The spatial and channel attention weights are combined to generate an integrated attention map:

[0144]

[0145] where denotes the broadcast element-wise multiplication, i.e., the spatial attention tensor spatial and the integrated attention tensor are extended to shape B x H x W x C through the broadcast mechanism and then respectively multiplied element-wise with the channel attention tensor of shape B x H x W x C and the aligned feature map scale , Conv 1×1 is a convolution with kernel size 1 and output channel number 1, BN is batch normalization, and Sigmoid denotes the Sigmoid activation function. Finally, the refined feature

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

[0147] After obtaining the refined features from the encoder stage through the spatial-channel dual attention mechanism , the MCFB integrates them with the current decoder feature. All cross-scale features are concatenated and processed through the LGF-VSS block, and finally combined with the original decoder feature through a residual connection, which can be represented as:

[0148]

[0149] where denote the different refined features obtained after the above feature scale alignment, spatial attention calculation, channel attention calculation, and integrated attention map generation and feature refinement of the outputs from the lower stage, the same stage, and the higher stage of the encoder, respectively, Concat(·) represents concatenating all refined features along the channel dimension, represents feeding the concatenated feature 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 perform the final spatial resolution upsampling and pixel-level classification on the feature map output by the final stage decoder to obtain the pixel-level mask segmentation result. The specific implementation steps are as follows:

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

[0153]

[0154] where Linear is a linear layer with an output channel number of 4 times the input channel number.

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

[0156]

[0157] where Reshape represents tensor reshaping, and is reshaped to

[0158] LN(·) represents layer normalization.

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

[0160]

[0161] where Conv 1×1 represents a convolution with a kernel of 1 and an output channel number of the class number of the target segmentation task, and Softmax represents the Softmax function, which converts the value at each pixel position into a probability distribution, indicating the probability that the pixel belongs to each class. Finally, the channel number with the maximum probability at each pixel position in Mask' is taken as the final predicted class index of the pixel, thereby generating a complete pixel-level semantic segmentation mask graph Mask.

[0162] Step 3, use the gradient statistics-based adaptive hierarchical loss function combined with certain training settings to train the LGFVM-UNet model on the training set, and optimize the network model parameters to make the network achieve optimal effect.

[0163] The gradient statistics adaptive hierarchical loss balancing strategy proposed by the application uses different weights for the losses of different decoder layers, and adaptively changes the learning dynamics throughout the training process, thereby achieving a better learning effect.

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

[0165] The gradient statistical adaptive hierarchical loss balancing strategy is implemented as follows:

[0166] (1) Gradient amplitude calculation:

[0167] For the lth decoder stage of the training epoch t, t≥1, first calculate the gradient amplitude:

[0168]

[0169] Where |θ| represents the total number of network parameters, is the loss of the lth decoder stage. This gradient amplitude reflects the current learning state and optimization potential of each supervised layer.

[0170] (2) Momentum smoothing weight update:

[0171] To ensure stable weight update and prevent violent oscillation, a momentum-based smoothing method is used to calculate the weight of the current decoder stage:

[0172]

[0173] Where α∈[0,1] is the momentum coefficient that controls the retention of historical information, exp represents the exponential function, and τ>0 is the temperature parameter that controls the sharpness of the weight distribution. If If there is no current training epoch 1, then The value is 0.

[0174] (3) Final composite loss:

[0175] The final composite loss in any epoch combines the adaptively weighted supervised loss and the main output loss to obtain:

[0176]

[0177] Where λ is the balancing hyperparameter that determines the relative importance of the main output loss, and The calculation formula is the same as that of The complete pixel-level semantic segmentation mask map generated by the low-resolution decoder is used as the final output, so the real segmentation mask map needs to be down-sampled to the same resolution as the final output for subsequent calculations, and In the implementation step, this operation is not required. In the implementation step, the final output of the low-resolution decoder is directly used as the final output of the model, and the real segmentation mask map is not down-sampled. The final output of the low-resolution decoder is directly used as the final output of the model, and the real segmentation mask map is not down-sampled. The final output of the low-resolution decoder is directly used as the final output of the model, and the real segmentation mask map is not down-sampled. The final output of the low-resolution decoder is directly used as the final output of the model, and the real segmentation mask map is not down-sampled. The final output of the low-resolution decoder is directly used as the final output of the model, and the real segmentation mask map is not down-sampled.

[0178]

[0179] where α and β are determined cross-entropy loss Dice similarity coefficient loss Balancing hyperparameters of relative importance, N is the total number of image pixels, C is the number of target segmentation task categories, y i,c is the one-hot encoding of the pixel on the real segmentation mask map or the real segmentation mask map after down-sampling to the same resolution as the final output, which is 1 only when the pixel belongs to the c-th class, otherwise it is 0, is the one-hot encoding of the pixel on the model predicted mask map, σ is the Softmax function, ε is a small smoothing coefficient, which is 0.000001 to avoid division by zero.

[0180] Based on the gradient statistical adaptive hierarchical loss balancing strategy proposed in the present application, the training method proposed in the present application is implemented by Python and uses the PyTorch deep learning framework. The experiment is performed on a high-performance server equipped with an NVIDIA GeForce RTX4090 GPU, and the operating system is Ubuntu 20.04. After running for a certain period of time and a certain number of epochs, the training is completed.

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

[0182] 1. Learning rate scheduling: In the training process, the cosine annealing learning rate scheduling strategy is adopted, the initial learning rate is set to 0.001, and the learning rate is gradually decreased from the initial value to the minimum learning rate set in Table 1 according to the cosine function within a certain number of epochs of the total training period.

[0183] 2. Early stopping strategy: If the performance on the validation set does not improve within 20 epochs, the training is stopped in advance to avoid overfitting.

[0184] 3. L2 regularization: Prevents model overfitting by L2 regularization (weight decay). The regularization parameter is adjusted step by step during the model training process to ensure the best regularization effect.

[0185] Step 4, test the performance of the resulting network model using the test set.

[0186] After training the LGFVM-UNet model on the test set divided by the five public medical image datasets mentioned in step 1, the performance of the LGFVM-UNet model in medical image segmentation is evaluated by different indicators. All the indicators involved in this embodiment are as follows:

[0187] 1. Dice Similarity Coefficient (DSC): used to measure the volume overlap between the prediction and the true label.

[0188] 2. 95% Hausdorff Distance (HD95): used to evaluate the boundary delineation accuracy.

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

[0190] The results of different experiments on the five public 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 exhibits excellent and stable performance in the evaluation of ISIC2017, ISIC2018, CVC-ClinicDB, Synapse and ACDC datasets, and the comprehensive indicators (such as DSC generally higher than 91%, class average DSC reaching 88.74% and 92.40%) and other key indicators (such as mIOU, ACC, Sen, Spe) are kept at a high level, fully verifying 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; on the ACDC dataset, the class average DSC reaches 92.40% and the HD95 is 21.97; on the ISIC2017 dataset, the DSC is 91.70% and the mIoU is 84.67%; on the ISIC2018 dataset, the DSC is 91.35% and the mIoU is 84.09%; on the CVC-ClinicDB dataset, the DSC is 91.03% and the mIoU is 83.55%. This shows that the LGFVM-UNet model can well perform medical image segmentation under different modalities and different regions, and has high diagnostic accuracy and reliability.

[0201] Through the above test evaluation, the LGFVM-UNet model proposed by the present application exhibits good segmentation performance in medical images and can realize automatic and intelligent medical image segmentation function. The high efficiency and accuracy of the model provide a reliable auxiliary diagnosis tool for clinical practice, and has wide application prospect and significant clinical value.

[0202] The above embodiments describe the technical solutions and advantages of the present application in detail. It should be understood that the above description is only a specific embodiment of the present application and is not intended to limit the present application. Any modification, supplement and equivalent replacement made within the principle range of the present application shall be included in the protection scope of the present application.

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. (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, In step (2), the LGF-VSS module includes the LGF-SS2D module and the dynamic gating fusion block; 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. 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. 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.

4. The 2D medical image segmentation method based on Mamba and UNet according to claim 3, 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.

5. The 2D medical image segmentation method based on Mamba and UNet according to claim 3, 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.

6. 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.

7. The 2D medical image segmentation method based on Mamba and UNet according to claim 1, characterized in that, In step (2), the MCFB module fuses the refined feature map with the feature maps output from all stages of the encoder. The specific process is as follows: The resolution of feature maps from all stages of the encoder output can be unified 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.

8. 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: Where |θ| represents the total number of network parameters, It is the loss of the decoder in the l-th stage excluding the last stage; The loss weights for each stage are updated using momentum smoothing: 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; The final loss is the sum of the weighted multi-level supervision loss and the main output loss: 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.

9. 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.

10. 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-9.

Images