MUNet medical image segmentation model and medical image segmentation method

Abstract

The invention discloses a medical image segmentation model and segmentation method, and the method comprises the steps: obtaining a training data set, and carrying out the preprocessing; secondly, constructing a medical image segmentation model; then, training a medical image segmentation model and a comparison model, and carrying out decoding performance comparison; and finally, inputting the obtained medical image into the trained medical image segmentation model to obtain a segmented medical image. A pure vision Mama (VMama)-based encoder-decoder structure is employed, and jumper connections are fused to reserve spatial information in networks of different scales. According to the design, a comprehensive feature learning process is promoted, and complex details and wider semantic contexts in the medical image can be captured.

Description

Technical Field

[0001] The present invention belongs to the field of biomedical engineering, and in particular relates to a MUNet medical image segmentation model and a medical image segmentation method. Background Art

[0002] In the field of medical image analysis, Convolutional Neural Network (CNN) has established a solid benchmark in many medical image processing tasks with its powerful local feature extraction capabilities. CNN scans the image layer by layer through convolutional layers, effectively capturing local detail features such as edges and textures in the image, which are essential for the preliminary identification and classification of abnormal areas in medical images. UNet is an important benchmark model in CNN-based medical image segmentation models. However, despite CNN's excellent performance in local feature extraction, its capabilities are relatively limited when dealing with complex tasks that require extensive contextual information or long-range dependencies, such as accurate segmentation of lesion areas or vascular networks.

[0003] At the same time, the rise of Vision Transformers (ViT) has brought a new perspective to medical image analysis. The ViT model achieves a deep understanding of the global context of the image by introducing the self-attention mechanism commonly used in natural language processing. The self-attention mechanism allows the model to consider the information of all other locations in the image when processing each pixel or feature, thereby capturing more complex spatial relationships and long-distance dependencies. This feature makes ViT show great potential in understanding the global structure in medical images and identifying complex pathological patterns.

[0004] However, both CNN and ViT face certain limitations when applied to medical image analysis alone. Although CNN is good at capturing local details, it may ignore the integration of global information; and although ViT can grasp the global context of medical impact, when processing high-resolution medical images, computational complexity and memory consumption may become bottlenecks, and it may focus too much on the global situation and ignore some key local details.

[0005] Recently, significant progress has been made in the field of state space models (SSM), especially the rise of structured SSM, which provides an efficient and powerful solution for processing long sequence data. On this basis, the Mamba model further enhances the performance of structured SSM through carefully designed selectivity and hardware optimization strategies. In order to extend this advantage to the field of computer vision, especially the field of medical image analysis, which requires high precision, the visual state space model (VMamba) came into being. The model innovatively introduces the cross-scan module (CSM), cleverly solves the problem of spatial domain traversal, and successfully transforms non-causal visual images into ordered map sequences, greatly enhancing the applicability and effectiveness of the Mamba model in complex visual tasks.

[0006] Therefore, by utilizing the advantages of both CNN and VMamba, a UNet (MUNet) based on the Mamba framework is constructed to achieve high-precision medical image segmentation. Summary of the invention

[0007] Due to the above-mentioned defects in the prior art, the present invention provides a method for realizing high-precision segmentation of medical images, so as to overcome the problems of poor precision, high computational complexity and large memory consumption of traditional medical image segmentation methods.

[0008] In order to achieve the above object, the present invention provides the following technical solutions:

[0009] A MUNet medical image segmentation model includes an encoder, a bottleneck, a decoder and skip connections; wherein the encoder, the bottleneck and the decoder are all built based on VMamba blocks, each stage of the encoder and the decoder and the bottleneck include 2 VMamba blocks, and the decoder also includes a block expansion layer; each layer of the encoder and the decoder uses skip connections.

[0010] The construction process of the encoder is as follows:

[0011] First, the feature dimension of an initial linear embedding layer is reduced to an arbitrary size, denoted as ;

[0012] Secondly, The dimensionally labeled input is passed through two consecutive VMamba blocks for feature learning, keeping the dimension and resolution unchanged;

[0013] Then, after three times of block merging as a downsampling process, the input is divided into four parts and concatenated;

[0014] Finally, layer normalization is performed after each block merging to regularize the dimensions, reduce the number of tokens by 1 / 2, and double the feature dimension.

[0015] The construction process of the decoder is as follows:

[0016] First, two consecutive VMamba blocks are used to reconstruct features; then, a block expansion layer is used to upsample the deep features, reducing the feature dimension while doubling the resolution.

[0017] Each layer of the encoder and decoder uses skip connections to fuse multi-scale features and the upgraded output, enhancing spatial details by merging shallow and deep layers; the subsequent linear layers maintain the dimensionality of the integrated feature set to keep it consistent with the upgraded resolution.

[0018] The VMamba block introduces a cross-scan module and integrates the convolution operation into the module; in the VMamba block, the input features first pass through a linear embedding layer and then are divided into two paths, one of which undergoes deep convolution and SiLU activation, then enters the SS2D module, and after post-layer normalization, is merged with the other branch that has undergone SiLU activation.

[0019] The medical image segmentation method comprises the following steps:

[0020] Step 1: Obtain the training data set and perform preprocessing;

[0021] Step 2: constructing the medical image segmentation model;

[0022] Step 3: training the medical image segmentation model and the comparison model, and comparing the decoding performance;

[0023] Step 4: Input the acquired medical image into the trained medical image segmentation model to obtain the segmented medical image.

[0024] In step 1, the specific process of obtaining the training data set and preprocessing is as follows:

[0025] Step 1.1, obtain the laboratory data set, according to the acquisition method of the standard data set ACDC MRI heart segmentation data set, obtain the laboratory data set;

[0026] Step 1.2: Preprocess the laboratory data set, including:

[0027] Step 1.2.1, image data processing, name the medical image slices in order and save them in a specified folder, traverse each slice in the folder, and save each slice as a JPEG format image;

[0028] Step 1.2.2, label data processing, traverse each slice in the label file, name them in order, save them in PNG format, and save them in the specified folder;

[0029] Step 1.2.3: Dataset division: divide the laboratory dataset into a training dataset and a test dataset according to the ratio of the ACDC MRI heart segmentation dataset.

[0030] In step 3, the specific process of training the medical image segmentation model and the comparison model is as follows:

[0031] Step 3.1, set model training related parameters;

[0032] Step 3.2: After each set number of iterations, evaluate the network performance on the validation set;

[0033] Step 3.3: Determine whether the network performance has reached the new best performance on the validation set. If so, save the model weights; otherwise, return to step 3.2.

[0034] In the specific process of re-evaluating the network performance, a wide range of evaluation indicators are used to evaluate MUNet and the benchmark methods. Among them, the similarity metrics include Dice coefficient, intersection over union ratio, accuracy, precision, sensitivity and specificity; the difference metrics include Hausdorff distance and average surface distance.

[0035] A computer-readable storage medium stores computer-readable instructions, which, when executed by a processor, call all or part of the steps of the method.

[0036] The above invention has the following advantages or beneficial effects:

[0037] The encoder-decoder structure based on pure visual Mamba (VMamba) is adopted, and skip connections are integrated to preserve spatial information in networks of different scales. This design promotes a comprehensive feature learning process that can capture complex details and broader semantic context in medical images. BRIEF DESCRIPTION OF THE DRAWINGS

[0038] The present invention and its features, appearance and advantages will become more apparent by reading the detailed description of non-limiting embodiments with reference to the following drawings. The same reference numerals indicate the same parts in all the drawings. The drawings are not drawn to scale, and emphasis is placed on illustrating the subject matter of the present invention.

[0039] Figure 1 This is a flow chart of the method for classifying motor imagery signals in Example 1.

[0040] Figure 2The structural diagram of UNet (MUNet) based on the Mamba framework constructed in Example 1.

[0041] Figure 3 This is a structural diagram of VMamba constructed in Example 1. DETAILED DESCRIPTION

[0042] The present invention is described in detail below in conjunction with the accompanying drawings and specific embodiments. This embodiment is implemented based on the technical solution of the present invention, and provides a detailed implementation method and specific operation process, but the protection scope is not limited to the following embodiments.

[0043] A MUNet medical image segmentation model includes an encoder, a bottleneck, a decoder and skip connections; wherein the encoder, the bottleneck and the decoder are all built based on VMamba blocks, each stage of the encoder and the decoder and the bottleneck include 2 VMamba blocks, and the decoder also includes a block expansion layer; each layer of the encoder and the decoder uses skip connections.

[0044] The medical image segmentation method comprises the following steps:

[0045] Step 1: Obtain the training data set and perform preprocessing;

[0046] Step 2: constructing the medical image segmentation model;

[0047] Step 3: training the medical image segmentation model and the comparison model, and comparing the decoding performance;

[0048] Step 4: Input the acquired medical image into the trained medical image segmentation model to obtain the segmented medical image.

[0049] Specific embodiments, such as Figures 1 to 3 As shown,

[0050] A medical image segmentation method comprises the following steps:

[0051] (1) Obtaining training data set:

[0052] 1. Collect data from subjects to form laboratory data sets, and preprocess the laboratory data sets:

[0053] 1.1: Obtain the laboratory dataset. The collection method of the laboratory dataset is exactly the same as the standard dataset ACDC (Automatic Cardiac Diagnosis Challenge) MRI heart segmentation dataset. The ACDC MRI heart segmentation dataset is the dataset used by the ACDC Challenge of MICCAI 2017, which aims to segment the left ventricle (LV), right ventricle (RV) and myocardium (Myo) in the diastolic (ED) and systolic (ES) frames of cardiac dynamic magnetic resonance imaging (cine-MRI). The dataset contains 150 cases, divided into 5 subclasses: NOR (normal), MINF (myocardial infarction with systolic heart failure), DCM (dilated cardiomyopathy), HCM (hypertrophic cardiomyopathy) and ARV (right ventricular abnormality), with 30 cases in each class. The above data is evenly divided into a training set of 100 cases and a test set of 50 cases.

[0054] 1.2: Preprocess the laboratory dataset:

[0055] 1.2.1: Image data processing, traverse each slice (two-dimensional layer in three-dimensional data) in the NIfTI file, and save each slice as a JPEG format image. The JPEG format is suitable for storing image data because it supports lossy compression but still maintains high visual quality. The slices should be named in sequence and saved in a specified folder.

[0056] 1.2.2: Label data processing, also traverse each slice in the NIfTI label file, since the label data usually contains integer category labels (such as 0 for background, 1-3 for different cardiac structures), it should be saved in PNG format to avoid the distortion of category labels that may be caused by JPEG compression. The slices should also be named in sequence and saved in the specified folder. The image data and label data are all adjusted to 224×224 size.

[0057] 1.2.3: Dataset division: The laboratory dataset is divided into training dataset and test dataset according to the ratio of ACDC MRI heart segmentation dataset.

[0058] (2) Construction of segmentation model:

[0059] The overall structure of the framework is as follows Figure 2 As shown in Figure 2, the architecture of MUNet consists of an encoder, a bottleneck, a decoder, and skip connections. The encoder, bottleneck, and decoder are all built based on VMamba blocks;

[0060] The 2D grayscale image size of the MUNet architecture input is , firstly it is split into small pieces similar to ViT, and then converted into a dimension of A one-dimensional sequence of . An initial linear embedding layer resizes the feature dimension to an arbitrary size, denoted as C. These small blocks of tokens are then processed through multiple VMamba blocks and block merging layers to create hierarchical features. The block merging layers handle downsampling and dimensionality increase, while the VMamba blocks focus on learning feature representations. The output resolution of each stage of the encoder is , , and The decoder consists of VMamba blocks and block expansion layers, following the encoder style so that the output feature sizes are exactly the same, thereby enhancing the spatial details lost in downsampling through skip connections. In each stage of the encoder and decoder, 2 VMamba blocks are used.

[0061] 2.1, build VMamba block:

[0062] The traditional state space model (SSM) is a linear time-invariant system function that uses hidden states Mapping , given As the evolution parameter, As state size The projection parameters and jump connections . The model can be expressed as a linear ordinary differential equation:

[0063]

[0064]

[0065] The discretized version of this linear model can be implemented using a zero-order hold transform, given a time scale parameter

[0066]

[0067]

[0068]

[0069]

[0070]

[0071] in . is approximated by a first-order Taylor series To refine. Figure 3As shown in Figure 1, the VMamba block further introduces a cross scan module (CSM) and then integrates the convolution operation into the module. In the VMamba block, the input features first pass through a linear embedding layer and then split into two paths. One branch undergoes depthwise convolution and SiLU activation, then enters the SS2D module and merges with the other branch after SiLU activation after post-layer normalization. Unlike typical ViT, this VMamba block abandons positional embedding and chooses a simplified structure without MLP, allowing denser modules to be stacked within the same depth budget.

[0072] 2.2, build the encoder:

[0073] In the encoder, after dimensionality reduction The dimensionally labeled input is passed through two consecutive VMamba blocks for feature learning, keeping the dimension and resolution unchanged. In the encoder of MUNet, three block merging is used as a downsampling process, which reduces the number of tokens by 1 / 2 and doubles the feature dimension by dividing the input into four parts, concatenating them, and then normalizing the dimension by layer normalization after each.

[0074] 2.3, build the decoder:

[0075] The decoder mimics the structure of the encoder, using two consecutive VMamba blocks for feature reconstruction and block expansion layers instead of pooling layers for upsampling deep features. These layers double the resolution while reducing the feature dimensionality.

[0076] 2.4, Bottleneck and Skip Connections:

[0077] Two VMamba blocks are used in the bottleneck of MUNet. Each layer of the encoder and decoder uses skip connections to fuse multi-scale features and upscaled outputs, enhancing spatial details by merging shallow and deep layers. The subsequent linear layer maintains the dimensionality of the integrated feature set to ensure consistency with the upscaled resolution.

[0078] (3) Train the decoding model and comparison model, and compare the decoding performance (UNet, Attention UNet, TransUNet and Swin-UNet):

[0079] The model training parameters are set as follows: Epoch is set to 10000, Batch_size is set to 24, and learning rate is set to , momentum was set to 0.9, and weight decay was set to 0.0001. After each 200 iterations, the network performance was evaluated on the validation set, and the model weights were only saved when a new best performance was achieved on the validation set.

[0080] A wide range of evaluation metrics are used to evaluate MUNet and the comparison methods. Similarity metrics include Dice coefficient, intersection over union (IoU), accuracy, precision, sensitivity, and specificity; difference metrics include Hausdorff distance (HD) and average surface distance (ASD). MUNet has achieved better performance in many metrics.

[0081] The medical image segmentation method of the present invention adopts an encoder-decoder structure based on pure visual Mamba (VMamba) and incorporates skip connections to preserve spatial information in networks of different scales. This design promotes the comprehensive feature learning process and can capture complex details and broader semantic context in medical images. The method can greatly improve the segmentation accuracy of medical images and has good application prospects.

[0082] The preferred embodiments of the present invention are described above. It should be understood by those skilled in the art that they can implement variations by combining the prior art and the above embodiments, which will not be described in detail here. Such variations do not affect the essential content of the present invention, and will not be described in detail here.

[0083] It should be understood that the present invention is not limited to the above-mentioned specific implementation modes, and the devices and structures not described in detail should be understood to be implemented in a common manner in the art; any technician familiar with the art can make many possible changes and modifications to the technical solutions of the present invention by using the above-mentioned methods and technical contents without departing from the scope of the technical solutions of the present invention, or modify them into equivalent embodiments of equivalent changes, which does not affect the essential contents of the present invention. Therefore, any simple modification, equivalent change and modification made to the above embodiments according to the technical essence of the present invention without departing from the contents of the technical solutions of the present invention still fall within the scope of protection of the technical solutions of the present invention.

Claims

1. A MUNet medical image segmentation model, characterized by: It includes an encoder, a bottleneck, a decoder and a skip connection; wherein the encoder, the bottleneck and the decoder are all built based on VMamba blocks, and each stage of the encoder and the decoder and the bottleneck include 2 VMamba blocks, and the decoder also includes a block extension layer; each layer of the encoder and the decoder adopts a skip connection.

2. The MUNet medical image segmentation model according to claim 1, characterized in that: The construction process of the encoder is as follows: First, the feature dimension of an initial linear embedding layer is reduced to an arbitrary size, denoted as ; Secondly, The dimensionally labeled input is passed through two consecutive VMamba blocks for feature learning, keeping the dimension and resolution unchanged; Then, after three times of block merging as a downsampling process, the input is divided into four parts and concatenated; Finally, layer normalization is performed after each block merging to regularize the dimensions, reduce the number of tokens by 1 / 2, and double the feature dimension.

3. The MUNet medical image segmentation model according to claim 1, characterized in that: The construction process of the decoder is as follows: First, two consecutive VMamba blocks are used to reconstruct features; then, a block expansion layer is used to upsample the deep features, reducing the feature dimension while doubling the resolution.

4. The MUNet medical image segmentation model according to claim 1, characterized in that: Each layer of the encoder and decoder uses skip connections to fuse multi-scale features and the upgraded output, enhancing spatial details by merging shallow and deep layers; the subsequent linear layers maintain the dimensionality of the integrated feature set to keep it consistent with the upgraded resolution.

5. The MUNet medical image segmentation model according to claim 1, characterized in that: The VMamba block introduces a cross-scan module and integrates the convolution operation into the module; in the VMamba block, the input features first pass through a linear embedding layer and then are divided into two paths, one of which undergoes deep convolution and SiLU activation, then enters the SS2D module, and after post-layer normalization, is merged with the other branch that has undergone SiLU activation.

6. A medical image segmentation method, characterized in that: The steps include: Step 1: Obtain the training data set and perform preprocessing; Step 2: construct the medical image segmentation model described in any one of claims 1 to 5; Step 3: training the medical image segmentation model and the comparison model, and comparing the decoding performance; Step 4: Input the acquired medical image into the trained medical image segmentation model to obtain the segmented medical image.

7. The medical image segmentation method according to claim 6, characterized in that: In step 1, the specific process of obtaining the training data set and preprocessing is as follows: Step 1.1, obtain the laboratory data set, according to the acquisition method of the standard data set ACDC MRI heart segmentation data set, obtain the laboratory data set; Step 1.2: Preprocess the laboratory data set, including: Step 1.2.1, image data processing, name the medical image slices in order and save them in a specified folder, traverse each slice in the folder, and save each slice as a JPEG format image; Step 1.2.2, label data processing, traverse each slice in the label file, name them in order, save them in PNG format, and save them in the specified folder; Step 1.2.3: Dataset division: divide the laboratory dataset into a training dataset and a test dataset according to the ratio of the ACDC MRI heart segmentation dataset.

8. The medical image segmentation method according to claim 6, characterized in that: In step 3, the specific process of training the medical image segmentation model and the comparison model is as follows: Step 3.1, set model training related parameters; Step 3.2: After each set number of iterations, evaluate the network performance on the validation set; Step 3.3: Determine whether the network performance has reached the new best performance on the validation set. If so, save the model weights; otherwise, return to step 3.

2.

9. The medical image segmentation method according to claim 8, characterized in that: In the specific process of re-evaluating the network performance, a wide range of evaluation indicators are used to evaluate MUNet and the benchmark methods. Among them, the similarity metrics include Dice coefficient, intersection over union ratio, accuracy, precision, sensitivity and specificity; the difference metrics include Hausdorff distance and average surface distance.

10. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores computer-readable instructions, and when the computer-readable instructions are executed by a processor, all or part of the steps of the method described in any one of claims 6 to 9 are called.

Images