A lightweight multi-scale medical image segmentation method based on MSF-Mamba

Abstract

This invention discloses a lightweight multi-scale image segmentation method based on Visual Mamba (VMamba). The method includes: 1) constructing an MSF-Mamba network model; 2) the MSF-Mamba network consists of a symmetrical U-shaped encoder-decoder structure; 3) the first three layers of the encoder are convolutional layers used to extract local image features; the last three layers are four-branch parallel visual state space layers (PVSS layers) used to extract global image features; 4) a boundary enhancement module (BFEB) is added to each layer of the encoder to improve the model's segmentation accuracy for blurred boundaries; 5) the skip connections between the encoder and decoder adopt an S-Mamba structure for multi-scale feature fusion. Experiments show that this model has excellent segmentation performance, with only 96K model parameters, making it a lightweight, high-precision segmentation model with high application and promotion value.

Description

Technical Field

[0001] This invention relates to medical image segmentation, specifically to a lightweight multi-scale image segmentation method based on MSF-Mamba, which employs a U-shaped encoder-decoder structure and a skip connection mechanism, belonging to the interdisciplinary application field of artificial intelligence and medical imaging. Background Technology

[0002] Deep learning excels in image segmentation tasks. In the field of medical image segmentation, higher demands are being placed on its segmentation accuracy, especially the segmentation boundaries, and lightweight model structures are being continuously explored. Existing deep learning models widely adopt convolutional neural networks (CNNs) and attention-based Transformer networks. CNNs are good at capturing local features but lack global information feature representation; Transformer networks perform well in learning long-distance feature representations but also incur high computational costs. Albert Gu and Tri Dao proposed the state-space model Mamba at the end of 2023, which, while maintaining performance comparable to Transformer, achieves linear computational complexity and linear memory usage for long sequences, solving the problem of the huge computational overhead of Transformer on long sequences. VisionMamba proposed by Zhu Lianghui et al. and VMamba proposed by the Chinese Academy of Sciences, Huawei, and Pengcheng Laboratory extend the application of Mamba from one-dimensional sequence data to two-dimensional image space. It is a visual Mamba model (VMamba) with a global receptive field and linear complexity, achieving higher prediction accuracy than Vision Transformers (ViT) with fewer model parameters, showing great potential in the direction of lightweight design. Summary of the Invention

[0003] The purpose of this invention is to achieve a high-precision, lightweight medical image segmentation method. The technical solution provided by this invention is as follows:

[0004] like Figure 1 As shown, this invention provides a lightweight medical image segmentation method based on visual Mamba, MSF-Mamba, the steps of which include:

[0005] 1) Build a 6-layer U-shaped symmetrical encoder-decoder network MSF-Mamba, such as... Figure 2 As shown, the network includes 6 convolutional layers (Conv blocks), 6 max pooling layers (Maxpool), 6 boundary extraction modules (BFEB), 6 four-channel parallel visual state-space layers (PVSS), and 1 multi-scale feature fusion skip connection module (S-Mamba).

[0006] 2) Construct a 1-3 layer encoder: Each layer consists of a Conv block, Maxpool, and BFEB. The Conv block is used to extract local shallow features, and the BFEB is used to enhance edge information. The BFEB structure is as follows: Figure 3 As shown, it consists of a graph convolution Vig Block, a boundary generator (residual block + sigmoid function), and a weighted fusion module Re-Weight. The BFEB module transforms boundary prior knowledge into a learnable attention mechanism and achieves end-to-end boundary optimization through a structured process.

[0007] 3) Construct a 4-6 layer encoder: each layer consists of PVSS, Maxpool, and BFEB. PVSS is used to extract deep features and correlate global information. The PVSS layer structure is as follows: Figure 4 As shown, the four-branch parallel visual state space (VSS) module processes C / 4 channel features respectively;

[0008] 4) Construct the S-Mamba skip connection module, with the following structure: Figure 5 As shown, multi-scale features from the encoders of each layer are fused, and then the fused feature map is restored to the original scale and passed to the decoders of each layer.

[0009] 5) Symmetrical to the encoder, the decoder consists of three PVSS layers and three Conv blocks. The current decoding layer receives the decoding output from the upper layer and the S-Mamba fusion features, and after completing the decoding, it is input to the next decoding layer; the last layer of the decoder outputs the segmentation result.

[0010] 6) During model training, cross-entropy loss (BCE), Dice loss, and level set loss are weighted and calculated as shown in equation (1):

[0011] (1)

[0012] in, For learnable parameters, , Let BCE cross-entropy loss and Dice loss function be the k-th iteration, respectively, and their calculation formulas are shown in (2) and (3):

[0013] (2)

[0014] (3)

[0015] Where N represents all pixels in the skin lesion image, This represents the true label of pixel i. This represents the predicted probability of the model at pixel 𝑖. C represents the number of categories in the classification task. and Representing the first The true label value of each pixel in category c and the model's predicted probability. The smaller the loss value of both, the closer the model's prediction is to the true distribution.

[0016] Let be the level set loss function, which is used to drive the segmentation boundary to evolve towards the true edge by minimizing energy. Since the skin lesion image segmentation task only needs to distinguish between the lesion area and normal skin, it belongs to a binary classification problem. Therefore, the level set loss function can be simplified as shown in equation (4):

[0017]

[0018] (4)

[0019] in, For real-world labeled images, This is a probability map of the segmentation results predicted by the model. and This represents the average pixel intensity of the lesion area and normal tissue in the original image; It is a smooth heaviside function that achieves a continuous transition from probability to boundary conditions. This represents the image domain.

[0020] Experiments have shown that the advantages of this invention are that the model has a significant advantage in segmentation accuracy and the number of model parameters is greatly reduced, which is conducive to the deployment of the model at the embedding end. Attached Figure Description

[0021] Figure 1 Flowchart of an entity relation extraction method based on large models and dynamic prompts

[0022] Figure 2 MSF-Mamba network architecture diagram

[0023] Figure 3 Boundary Feature Enhancement Module (BFEB) Structure Diagram

[0024] Figure 4 PVSS module structure diagram

[0025] Figure 5 S-Mamba module structure diagram

[0026] Figure 6 Visualization of experimental results of different models on the ISIC2018 dataset Detailed Implementation

[0027] This invention discloses a lightweight multi-scale medical image segmentation method based on MSF-Mamba. The method includes: 1) constructing a 6-layer U-shaped symmetric encoder-decoder network structure, with three convolutional layers (Conv) and three PVSS layers; 2) adding a boundary feature enhancement module (BFEB) to each layer of the encoder; 3) the PVSS layers adopting a four-branch parallel structure based on VSS; 4) constructing a multi-scale fusion module (S-Mamba) between the encoder and decoder; and 5) constructing a composite loss function for model training. The technical solution of this invention is clearly and completely described below with reference to the accompanying figures and experiments using the publicly available dermoscopy datasets ISIC2017 and ISIC2018, as well as the Synapse multi-organ dataset.

[0028] 1. Construct a shallow feature extraction and boundary enhancement module. It consists of three layers: a convolutional block (Conv), a max-pooling layer, and a BFEB boundary enhancement module, as shown below. Figure 2 As shown, the specific steps are as follows:

[0029] 1.1 The first encoder layer consists of a Conv block convolutional layer, a Maxpool pooling layer, and a BFEB boundary enhancement layer. The medical image input is 224×224×3. The first convolutional layer, Conv2d, has a kernel size of 3×3, a stride of 1, and 8 channels. Then, it passes through the Maxpool pooling layer with a kernel size of 2×2 and a stride of 2×2, resulting in an output of 224×224×8.

[0030] 1.2 Construct the Boundary Enhancement Module (BFEB), such as... Figure 3 As shown, the algorithm consists of a graph convolutional Vig Block, a boundary generator (residual block + sigmoid function), and a weighted fusion module Re-Weight. After pooling, the features are processed by the graph convolutional Vig Block, which aggregates global and local contextual information by constructing node adjacency relationships, effectively capturing the topological structure of irregular shapes and outputting enhanced features. The boundary generator gradually refines the boundary features through two layers of dilated convolutions and a non-linear activation function (sigmoid). Re-Weight introduces a learnable weight matrix W for dynamic adjustment, ultimately outputting enhanced features. BFEB operates after the pooling layer of each encoder layer.

[0031] 1.3 Construct the second encoder layer, a convolutional layer (Conv block) with a kernel size of 3×3×8 and a stride of 1; a second max-pooling layer (Maxpool) with a kernel size of 2×2 and a stride of 2×2, outputting 112×112×16; the BFEB structure is as follows. Figure 3 ;

[0032] 1.4 Construct the third encoder layer, a convolutional layer (Conv block) with a kernel size of 3×3×24, a stride of 1, and 16 input channels; use a max pooling layer with a kernel size of 2×2 and a stride of 2×2; the BFEB layer has the same structure as above, with 56×56×24 output channels.

[0033] 2. Construct a deep feature extraction and boundary enhancement module, consisting of 3 layers. Each layer comprises a parallel visual state space module (PVSS), a pooling layer, and a BFEB layer. The specific steps are as follows;

[0034] 2.1 PVSS module, as shown Figure 4 As shown, the input feature X (number of channels C) is first processed by the linear normalization layer LayerNorm, and then divided into four channels, each channel having C / 4 sub-features. The sub-features of each channel are spatially modeled by the Visual State Space (VSS) module, and their outputs are fused with learnable weight coefficients through residual concatenation, and then concatenated along the channel dimension to obtain the original number of channels C. Finally, a second-level normalization and linear projection are performed for output. The calculation process is as shown in equation (1):

[0035]

[0036]

[0037] (1)

[0038]

[0039] 2.2 Construct the fourth encoder layer, with PVSS input of 28×28×24, each branch processes 6 channels of features, and output of 28×28×32; use Maxpool layer with a pooling kernel size of 2×2 and a stride of 2×2; the BFEB layer structure is the same as above, with an output of 14×14×32.

[0040] 2.3 Construct the fifth encoder layer. The PVSS input channel count is 14×14×32. Each branch processes 8 channels of features, and the output is 14×14×48. Maxpool is used with a pooling kernel size of 2×2 and a stride of 2×2. The BFEB layer structure is the same as above, and the output is 7×7×48.

[0041] 2.4 Construct a sixth encoder layer with 48 PVSS input channels. Each branch processes 12 channels of features, and the output is 7×7×64. This layer is saved as a bottleneck and does not participate in multi-scale fusion.

[0042] 3. A multi-scale fusion module, S-Mamba, is built between the encoder and decoder to fuse multi-scale features, such as... Figure 5 As shown, the specific steps are as follows:

[0043] 3.1 Save the output of each encoder layer:

[0044] t1 = out # b,8,H / 2,W / 2 (First layer)

[0045] t2 = out # b,16,H / 4,W / 4 (Second layer)

[0046] t3 = out # b,24,H / 8,W / 8 (Third layer)

[0047] t4 = out # b,32,H / 16,W / 16 (Fourth layer)

[0048] t5 = out # b,48,H / 32,W / 32 (Fifth layer)

[0049] 3.2 Concatenate all encoder features into a single feature map F_t3.

[0050] 3.2.1 After downsampling t1 (112×112×8), it is concatenated with t2 (56×56×16) to output t2 (56×56×24); then downsampled to obtain feature F_t2 (28×28×24);

[0051] 3.2.2 Upsample t5 (7×7×48) and concatenate it with t4 (14×14×32) to obtain t4 (14×14×80); Upsample the concatenated t4 to obtain F_t4 (28×28×80).

[0052] 3.2.3 Concatenation of F_t2, F_t4 and t3: F_t2(28×28×24) + t3(28×28×24) + F_t4(28×28×80) = F_t3(28×28×128); Linear projection of feature F_t3 to output feature F_t4;

[0053] 3.3 Process feature F_t4 using the SiLU activation function to output F_t5;

[0054] 3.4 After performing a one-dimensional convolutional layer on feature F_t4, the SiLU activation function and SSM layer are used for processing to output F_t6;

[0055] 3.5 Combine the features F_t5 and F_t6 from the two branches using the Hadamard product, perform linear projection to output the fused feature F, and calculate the process as shown in equation (2):

[0056] (2)

[0057] 3.6 Adjust the five feature maps output by Mamba to the same spatial size as the original encoder features t1-t5, replace the original t1-t5, and output them to each layer of the decoder respectively;

[0058] 4. Construct the decoder layer, which consists of 6 layers. The specific steps are as follows:

[0059] 4.1 decoder1: Build the PVSS module, input 7x7x64, output 7x7x48, add it to the new t5 (7x7x48) to get out5 (7x7x48);

[0060] 4.2 decoder2: Upsample out5 by 2 times, output 14x14x48, pass through the PVSS module, output 14x14x32, add it to the new t4 (14x14x32) to get out4 (14x14x32).

[0061] 4.3 decoder3: Upsample out4 by 2 times, output 28x28x32, pass through the PVSS module, output 28x28x24, add it to the new t3 (28x28x24) to get out3 (28x28x24).

[0062] 4.4 decoder4: Upsample out3 by a factor of 2 to output 56x56x24, pass it through the convolutional layer Conv2d(24,16,3x3) to output 56x56x16, add it to the new t2 (56x56x16) to get out2 (56x56x16).

[0063] 4.5 decoder5: Upsample out2 by a factor of 2 to output 112x112x16, output convolutional layer Conv2d(16,8,3x3) output 112x112x8, add it to the new t1 (112x112x8) to get out1 (112x112x8).

[0064] 4.6 decoder6: Upsamples out1 by a factor of 2 to output 224x224x8, outputs 224x224x1 through 1x1 convolution, and then activates the output through the sigmoid function.

[0065] 5. Input the dataset and construct a composite loss function. The loss function is calculated according to formula (3):

[0066] (3)

[0067] in, For learnable parameters, , Let BCE cross-entropy loss and Dice loss function be the values ​​for the k-th iteration, respectively. The level set loss function is introduced as an auxiliary loss function in the first iteration to improve supervision at the lesion boundary.

[0068] To demonstrate the effectiveness and lightweight nature of the network model in medical image segmentation, it was applied to a dermoscopy segmentation task for fair evaluation.

[0069] 5.1 The model training uses the SGD optimization function, and the SGD optimizer parameters are set as follows:

[0070] The learning rate Ir is set to the default value of 1e. -5 Momentum is 0.9; weight decay is 1e-4; minimum learning rate is 1e-5; batch size is 8.

[0071] 5.2 Before training, the dataset needs to be processed:

[0072] The ISIC2018 dataset contains 2594 dermoscopy images: 1816 for training, 260 for validation, and 518 for testing. The input images were uniformly cropped to a size of 256×256 pixels.

[0073] The model of this invention was trained on the ISIC2018 dataset and compared with methods such as the classic UNet (40.96 million parameters) and VM-UNet (27.42M) based on the Mamba architecture. All model experiments were run on a single NVIDIA RTX3090 GPU with 24GB of VRAM. The results are shown in Table 1.

[0074] Table 1. Comparative experimental results of different models on the ISIC2018 dataset.

[0075] Model Name mIoU DSC ACC SP SE Params GFLOPs U-Net 0.7786 0.8755 0.9405 0.9669 0.8586 2.009M 3.224 UNet++ 0.7831 0.8783 0.9402 0.9575 0.8865 23.01M 1.567 Transfuse 0.8063 0.8927 0.9466 0.9574 0.9128 26.30M 9.181 VM-UNet 0.8036 0.8911 0.9470 0.9653 0.8903 27.43M 4.112 UltraLight VM-UNet 0.8111 0.8957 0.9577 0.9747 0.8915 0.049M 0.060 MSF-Mamba(Ours) 0.8289 0.9064 0.9628 0.9823 0.8918 0.092M 0.088

[0076] Under the same experimental conditions, compared with the classic UNet, MSF-Mamba reduced the number of parameters by 21.8 times to only 0.092M, while improving mIoU and DSC by 5.03% and 3.09%, respectively; compared with VM-UNet based on the Mamba architecture, mIoU and DSC improved by 2.53% and 1.53%, respectively, while reducing the number of parameters by 298 times; compared with the lightweight UltraLight VM-UNet, the number of model parameters increased by nearly 2 times, but mIoU and DSC performance improved by 1.78% and 1.07%, respectively. Figure 6 As shown, the experimental results of different models on the ISIC2018 dataset are visualized.

[0077] To demonstrate the effectiveness of the PVSS, S-Mamba, and BFEB modules in improving segmentation performance of the MSF-Mamba model, a baseline model was set up. The first three stages of the encoder used convolutional layers, and the last three stages used the VSS module. Using its segmentation results as a baseline for comparison, the effectiveness of each module was verified on the publicly available ISIC2018 dermoscopy image segmentation task. The evaluation results are shown in Table 2.

[0078] Table 2. Experimental results of ablation of PVSS, S-Mamba, and BFEB modules.

[0079] ablation experiment PVSS S-Mamba BFEB Loss mIoU DSC ACC SP SE Baseline 0.2839 0.7981 0.8846 0.9495 0.9707 0.8861 +PVSS √ 0.2368 0.8221 0.9024 0.9608 0.9786 0.8910 +PVSS+S-Mamba √ √ 0.2323 0.8289 0.9064 0.9628 0.9823 0.8863 MSF-Mamba √ √ √ 0.2227 0.8379 0.9138 0.9654 0.9803 0.9050

[0080] Introducing the PVSS layer into the encoder improves performance metrics mIoU and DSC by 2.4% and 1.78%, respectively. Introducing the BFEB module further enhances mIoU and DSC by 0.9% and 0.74%, respectively, and strengthens the model's boundary feature extraction capabilities, demonstrating the effectiveness of the BFEB module. Introducing the S-Mamba mechanism into the skip connection part improves mIoU and DSC by 0.68% and 0.40%, respectively, proving the effectiveness of the S-Mamba module.

[0081] This invention implements a complete medical image segmentation method for a specific field. It constructs an image segmentation model based on a Mamba U-shaped encoder-decoder structure, introduces a four-channel parallel visual state space module (PVSS), a Mamba-based skip connection module (S-Mamba), and a boundary enhancement module (BFEB), realizes the fusion of multi-scale features and the enhancement of boundary features, and at the same time greatly reduces the number of model parameters, improves model efficiency and segmentation accuracy, and has broad application and promotion value.

[0082] Finally, it should be noted that the purpose of disclosing the embodiments is to help further understand the present invention. However, those skilled in the art will understand that various substitutions and modifications are possible without departing from the spirit and scope of the present invention and the appended claims. Therefore, the present invention should not be limited to the content disclosed in the embodiments, and the scope of protection of the present invention is defined by the claims.

Claims

1. A lightweight multi-scale medical image segmentation method based on MSF-Mamba, comprising the following steps: 1) Construct the first 3 layers of the encoder. Each layer consists of a Convolutional CNN layer (Conv), a Maxpooling layer, and a BFEB (Boundary Enhancement Module) layer. The convolutional kernel size of each layer is 3×3 with a stride of 1, and the Maxpooling kernel size is 2×2 with a stride of 2×2, completing the first stage of shallow feature extraction and boundary feature enhancement. The BFEB consists of a Graph Convolutional Vig Block, a boundary generator (residual block + sigmoid function), and a weighted fusion module (Re-Weight). 2) Construct the last 3 layers of the encoder. Each layer consists of a parallel visual state space module (PVSS), a max pooling layer, and a BFEB layer to complete deep feature extraction and boundary feature enhancement. 3) A skip connection module, S-Mamba, is built between the encoder and decoder for multi-scale feature fusion; 4) Construct the first 3 layers of the decoder. Each layer consists of PVSS and an upsampling layer. The input is a multi-scale feature splicing of the output of the previous layer decoder or encoder and the output of the skip connection. After being processed in parallel by the 4 channels of the PVSS module, it is upsampled and output to the next layer decoder. 5) Construct the last 3 layers of the decoder. Each layer consists of a convolutional layer (Conv) and an upsampling layer. The multi-scale features of the previous decoder input and the skip connection output are concatenated, then input into the Conv convolution and upsampled to complete the decoding output. 6) Input the dataset and train the segmentation task using a composite loss function. First, after inputting the image, the model calculates the loss function according to formula (1): （1） in, For learnable parameters, , Let BCE cross-entropy loss and Dice loss function be the values ​​for the k-th iteration, respectively. This is the level set loss function.

2. As described in claim 1, characterized in that, The parallel visual state space module PVSS in step 2) consists of a linear normalization layer LayerNorm, a visual state space module VSS, a linear normalization layer, and a linear projection layer. The calculation process is shown in equation (2), and its specific operation is as follows: (2) 1) The input feature X (number of channels C) is processed by the linear normalization layer LayerNorm, and then divided into four channels, each channel having C / 4 sub-features; 2) The features of the four parallel channels are spatially modeled by the Visual State Space (VSS) module, and then fused with learnable weight coefficients through residual concatenation; 3) Concatenate along the channel dimension to obtain the original number of channels C; 4) Perform secondary layer normalization and linear projection output.

3. As described in claim 1, characterized in that, In step 3), the skip connection module S-Mamba fuses the different scale features output from each encoder layer and then inputs them to each decoder layer. The specific operation is as follows: 1) Save the output of each encoder layer: t1 = out # b,8,H / 2,W / 2 (first layer); t2 = out # b,16,H / 4,W / 4 (second layer); t3 = out # b,24,H / 8,W / 8 (third layer); t4 = out # b,32,H / 16,W / 16 (Fourth layer); t5 = out # b,48,H / 32,W / 32 (fifth layer); 2) Concatenate all encoder features into a single feature map F_t3. The specific steps are as follows: After downsampling t1 (112×112×8) by 2, it is concatenated with t2 (56×56×16) to output t2 (56×56×24). Further downsampling yields feature F_t2 (28×28×24). Upsampling t5 (7×7×48) yields t5 (14×14×48), which is concatenated with t4 (14×14×32) to output t4 (14×14×80); further upsampling yields feature F_t4 (28×28×80). t3, F_t2, F_t4 concatenated: F_t2(28×28×24) + t3(28×28×24) + F_t4(28×28×80), output feature F_t3 (28×28×128); 3) Perform a linear projection on feature F_t3 and output feature F_t4; 4) Process feature F_t4 using the SiLU activation function and output F_t5; 5) After performing a one-dimensional convolutional layer on feature F_t4, process it using the SiLU activation function and an SSM layer to output F_t6; 6) Combine the features F_t5 and F_t6 of the two branches through the Hadamard product, perform linear projection to output the fused feature F, and calculate the process as shown in equation (3). (3) The fused feature F is input to each layer decoder.

4. As described in claim 1, characterized in that, In step 6), the composite loss function is calculated according to formula (4): （4） in For learnable parameters, , These are the BCE cross-entropy loss and Dice loss function for the kth iteration, respectively. The level set loss function is introduced as an auxiliary loss function in the first iteration to enhance supervision at the lesion boundary. Since the skin lesion image segmentation task only needs to distinguish between the lesion area and normal skin, it is a binary classification problem. Therefore, the level set loss function can be simplified as shown in equation (5): (5) Where G is the ground truth labeled image, and P is the probability map of the segmentation result predicted by the model. and H represents the average pixel intensity of the lesion area and normal tissue in the original image; ∗ It is a smooth Heaviside function that achieves a continuous transition from probability to boundary, and Ω represents the image domain.

Images