An image segmentation method and device based on the SM-UNet image segmentation model
By using the SM-UNet image segmentation model, combined with the sliding window mechanism and the SwinVSS Block, the problems of high computational complexity and insufficient local feature capture of Transformer in medical image segmentation are solved, and effective extraction and accurate segmentation of multi-scale features are achieved.
Patent Information
- Application Number
- CN202510254398.X
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2025-03-05
- Publication Date
- 2025-11-14
- Estimated Expiration
- 2045-03-05
AI Technical Summary
Existing Transformer-based medical image segmentation models suffer from high computational complexity, inability to effectively capture local features, and insufficient understanding of multi-size features.
The SM-UNet image segmentation model is adopted, which combines convolutional modules, encoders and decoders. Using a sliding window mechanism and SwinVSS blocks, feature extraction and merging are performed through convolutional operations and patch embedding segmentation layers to capture multi-scale features.
It effectively reduces computational complexity, improves the ability to capture local features, adapts to the needs of features of different sizes, and enhances the accuracy and efficiency of medical image segmentation.
Smart Images

Figure CN120088269B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of image segmentation, and in particular to an image segmentation method and device based on the SM-UNet image segmentation model. Background Technology
[0002] Existing solutions for medical image segmentation are mostly based on TransUNet. TransUNet is a deep learning model that combines the Transformer and UNet architectures, specifically designed for medical image segmentation. Transformer has achieved great success in Natural Language Processing (NLP) and, with technological advancements, has been attempted to be applied to computer vision. However, directly applying Transformer to computer vision tasks has revealed some shortcomings:
[0003] (i) High computational complexity
[0004] 1. Long sequences lead to slow computation.
[0005] When using Transformers for computer vision tasks, images are typically divided into small patches. For example, for high-resolution images, further subdividing patches to capture more features can result in very long sequences. Taking high-resolution images (such as 1000×1000 or 800×800) as an example, a patch size of 16×16 can generate a large number of patches, leading to sequences that can reach lengths of thousands.
[0006] 2. The computational complexity of self-attention mechanisms
[0007] The computational complexity of the self-attention mechanism is O(n^2). 2 ), where n is the length of the sequence. When processing images, due to the long sequence length, the computational complexity of the self-attention mechanism increases quadratically with the increase of image size. For example, when the image size doubles, the sequence length may double, and the computational cost will become four times the original.
[0008] (ii) Inability to effectively capture local features
[0009] Global Relationship Processing vs. Locality of Visual Information: Much visual information in an image depends on local relationships, such as textures and edges within an object. Traditional Transformers handle global relationships, performing self-attention computation on the entire input sequence, neglecting this local information characteristic. In Convolutional Neural Networks (CNNs), the convolution kernel slides across local regions to perform convolution operations, effectively capturing local features. Transformers lack this locality-based operation.
[0010] (iii) Insufficient understanding of multi-size characteristics
[0011] Single-scale feature extraction, as seen in early Transformer vision models like ViT (Vision Transformer), employs a single scale, processing features at a consistently 16×16 patch size, with the feature dimension remaining constant throughout the network. However, multi-scale features are crucial in many downstream tasks of vision tasks, such as object detection and image segmentation. In object detection, objects of different sizes require features of different scales for accurate identification; in image segmentation, different regions may require features of different resolutions for precise boundary delineation. Single-scale feature extraction performs poorly in these tasks because it cannot accommodate the multi-scale feature requirements of different objects and scenes.
[0012] Therefore, there is a need for image segmentation methods and devices that are less complex and capable of extracting features at multiple scales and comprehensively. Summary of the Invention
[0013] The purpose of this invention is to overcome the above-mentioned shortcomings in the prior art and to provide an image segmentation method and device based on the SM-UNet image segmentation model.
[0014] To achieve the above-mentioned objectives, the present invention provides the following technical solution:
[0015] An image segmentation method based on the SM-UNet image segmentation model includes the following steps:
[0016] Obtain the image to be segmented;
[0017] The image to be segmented is input into a pre-built SM-UNet image segmentation model;
[0018] The SM-UNet image segmentation model outputs image segmentation results;
[0019] The SM-UNet image segmentation model includes a convolutional module, an encoder, a skip connection module, and a decoder.
[0020] The convolution module is used to perform convolutional segmentation processing on the image to be segmented and output several feature map patches.
[0021] The encoder includes two branches for sliding window analysis, used to progressively extract and compress features from several feature map patches.
[0022] The jump connection module is used to retrieve spatial details lost during the process;
[0023] The decoder includes two branches that perform sliding window analysis to gradually restore the resolution of the feature map, perform detail reconstruction, and output image segmentation results.
[0024] As a preferred embodiment of the present invention, the convolution module includes a convolutional layer and a patch embedding segmentation layer;
[0025] The convolutional layer is used to convolve the image to be segmented into a feature map;
[0026] The Patch Embedding segmentation layer is used to segment the feature map into several feature map patches of the same size.
[0027] As a preferred embodiment of the present invention, the encoder includes a plurality of SwinVSS Blocks and a PatchMerging layer;
[0028] The SwinVSS Block is used to capture global and local features;
[0029] The Patch Merging layer is used to merge adjacent feature map patches.
[0030] As a preferred embodiment of the present invention, the SwinVSS Block includes a first VSS Block and a second VSS Block connected in sequence;
[0031] The first VSS block includes a normalized layer and a fully connected layer connected in sequence; wherein, a first branch and a second branch are provided between the normalized layer and the fully connected layer; the first branch includes a fully connected layer, a depth-separable convolutional layer, a W-2D selective scan layer, and a normalized layer connected in sequence; the second branch includes a fully connected layer;
[0032] The second VSS block includes a normalized layer and a fully connected layer connected in sequence; wherein, a third branch and a fourth branch are provided between the normalized layer and the fully connected layer; the third branch includes a fully connected layer, a depth-separable convolutional layer, an SW-2D selective scanning layer, and a normalized layer connected in sequence; the fourth branch includes a fully connected layer.
[0033] As a preferred embodiment of the present invention, the W-2D selective scanning layer is used to divide the feature map patch that makes up the image to be segmented into M*M calculation windows;
[0034] The SW-2D selective scanning layer is used to slide the calculation window of the W-2D selective scanning layer by 1 / 2 window unit length in both the height and width dimensions.
[0035] As a preferred embodiment of the present invention, the output expression of the W-2D selective scanning layer in the first VSS block is:
[0036] z l =W-SS2D(LN(z) l-1 ))+z l-1 ,
[0037] Among them, z l For the output of the W-2D selective scan layer, z l-1 SS2D() is the input for the W-2D selective scan layer, where W is the initial window selection, SS2D() is the 2D selective scan, and LN() is the normalization process.
[0038] The output expression of the SW-2D selective scan layer in the second VSS block is:
[0039] Z l+1 =SW-SS2D(LN(z) l ))+z l ,
[0040] Among them, z l+1 This is the output of the SW-2D selective scan layer, where SW is a sliding window processing layer.
[0041] As a preferred embodiment of the present invention, the depth-separable convolutional layer includes 8 convolutional layers, and the number of channels in each layer is [16, 32, 64, 128, 128, 64, 32, 16].
[0042] As a preferred embodiment of the present invention, the decoder includes several Patch Expanding layers and a SwinVSS Block;
[0043] The Patch Expanding layer is used to restore the resolution of the feature map and increase the number of feature map patches.
[0044] As a preferred embodiment of the present invention, the jump connection module includes two SwingVSS blocks.
[0045] An image segmentation device based on the SM-UNet image segmentation model includes at least one processor and a memory communicatively connected to the at least one processor; the memory stores instructions executable by the at least one processor, which, when executed by the at least one processor, enable the at least one processor to perform any of the methods described above.
[0046] Compared with the prior art, the beneficial effects of the present invention are as follows:
[0047] This invention proposes a novel encoder-decoder Mamba framework, which applies a sliding window to the Mamba-based UNet segmentation model to establish the SM-UNet image segmentation model, realizing Swin Mamba UNet (SMUNet) for automatic medical image segmentation. It mainly combines the advantages of Mamba's linear complexity and the cross-window interaction capability brought by the sliding window to effectively optimize the structure of the standard U-shaped architecture. Attached Figure Description
[0048] Figure 1 This is a flowchart illustrating an image segmentation method based on the SM-UNet image segmentation model as described in Embodiment 1 of the present invention.
[0049] Figure 2 This is a schematic diagram of the structure of the SM-UNet image segmentation model in the image segmentation method based on the SM-UNet image segmentation model described in Embodiment 2 of the present invention;
[0050] Figure 3 This is a schematic diagram of the SwinVSS Block structure in the image segmentation method based on the SM-UNet image segmentation model described in Embodiment 2 of the present invention;
[0051] Figure 4 This is a schematic diagram illustrating the effect of the shift window method in the image segmentation method based on the SM-UNet image segmentation model described in Embodiment 2 of the present invention.
[0052] Figure 5 This is a schematic diagram of the workflow of the SS2D module in the image segmentation method based on the SM-UNet image segmentation model described in Embodiment 2 of the present invention;
[0053] Figure 6 This is a schematic diagram of the structure of an image segmentation device based on the SM-UNet image segmentation model, which utilizes the image segmentation method based on the SM-UNet image segmentation model described in Embodiment 1, as described in Embodiment 4 of the present invention. Detailed Implementation
[0054] The present invention will be further described in detail below with reference to experimental examples and specific embodiments. However, this should not be construed as limiting the scope of the above-mentioned subject matter of the present invention to the following embodiments; all technologies implemented based on the content of the present invention fall within the scope of the present invention.
[0055] TransUNet is a deep learning model that combines the Transformer and UNet architectures, specifically designed for medical image segmentation. The following is a detailed explanation of how TransUNet achieves medical image segmentation:
[0056] I. Components and constituent parts
[0057] Components:
[0058] Transformer Layer: Used to process sequential data and capture long-distance dependencies.
[0059] CNN (Convolutional Neural Network): Used to extract local features.
[0060] MLP (Multilayer Perceptron): Used for nonlinear transformations.
[0061] LayerNorm (layer normalization): Used to stabilize the training process.
[0062] MSA (Multi-head Self-Attention Mechanism): Used to capture the relationships between different parts.
[0063] Constituent components:
[0064] Embedded sequence: The sequence obtained after preprocessing the input medical image.
[0065] Transformer block: Contains multiple Transformer Layers.
[0066] A CNN block contains multiple convolutional layers.
[0067] Upsampling and downsampling modules: used to adjust the resolution of the feature map.
[0068] Segmentation head: Used to generate the final segmentation result.
[0069] II. Connection relationships between components
[0070] a) The embedded sequence is first normalized using LayerNorm.
[0071] b) The normalized sequence enters the MSA and captures global information through a multi-head self-attention mechanism.
[0072] c) The features after MSA are further processed by LayerNorm and MLP.
[0073] d) The processed features are transformed into a shape suitable for CNN processing through a reshape operation.
[0074] e) CNN blocks perform convolution operations on features to extract local features.
[0075] f) The resolution of the feature map is gradually reduced by using a downsampling module to increase the receptive field.
[0076] g) During the downsampling process, the feature map fuses features at different scales through feature concatenation.
[0077] h) The resolution of the feature map is gradually restored through the upsampling module.
[0078] i) Finally, the feature map generates the segmentation result through the segmentation head.
[0079] III. Implementation and Function of Each Unit
[0080] Transformer Layer: Captures global dependencies through a self-attention mechanism, enhancing the model's ability to process long-distance information.
[0081] CNN: Extracts local features and captures detailed information in an image through convolution operations.
[0082] MLP: Introducing nonlinear transformations to enhance the expressive power of the model.
[0083] LayerNorm: Stabilizes the training process and accelerates convergence.
[0084] MSA: Through a multi-head self-attention mechanism, it captures the relationships between different parts, enhancing the model's global perception capability.
[0085] Downsampling: Reduce the resolution of the feature map by pooling or convolution operations to increase the receptive field.
[0086] Upsampling: Restores the resolution of the feature map through deconvolution or interpolation operations to generate high-resolution segmentation results.
[0087] Segmentation Header: Typically consists of convolutional layers and activation functions, generating the final segmentation mask.
[0088] IV. Implementation process of the program algorithm
[0089] 1. Data preprocessing: Performing preprocessing operations such as normalization, cropping, and scaling on medical images.
[0090] 2. Embedded sequence generation: Convert the preprocessed image into a sequence form and use it as input to the model.
[0091] 3. Transformer block processing:
[0092] i. Perform layer normalization on the embedded sequence.
[0093] ii. Capture global information via MSA.
[0094] iii. Perform layer normalization and MLP processing again.
[0095] 4. Feature reshape: Converts the features output by the Transformer block into a shape suitable for CNN processing.
[0096] 5. CNN block processing:
[0097] iv. Extract local features through convolutional layers.
[0098] v. Reduce the feature map resolution using a downsampling module.
[0099] vi. To fuse features at different scales through feature concatenation.
[0100] 6. Upsampling module processing:
[0101] vii. Gradually restore the resolution of the feature map through upsampling operations.
[0102] 7. The segmentation head generates the segmentation results:
[0103] viii. Generate the final segmentation mask through convolutional layers and activation functions.
[0104] 8. Loss Calculation and Optimization: Calculate the loss between the segmentation result and the true label, and optimize the model parameters through backpropagation.
[0105] Through the steps described above, TransUNet effectively combines the advantages of Transformer and CNN to achieve high-quality medical image segmentation.
[0106] Example 1
[0107] like Figure 1 As shown, an image segmentation method based on the SM-UNet image segmentation model includes the following steps:
[0108] Obtain the image to be segmented;
[0109] The image to be segmented is input into a pre-built SM-UNet image segmentation model;
[0110] The SM-UNet image segmentation model outputs image segmentation results;
[0111] The SM-UNet image segmentation model includes a convolutional module, an encoder, a skip connection module, and a decoder.
[0112] The convolution module is used to perform convolutional segmentation processing on the image to be segmented and output several feature map patches.
[0113] The encoder includes two branches for sliding window analysis, used to progressively extract and compress features from several feature map patches.
[0114] The jump connection module is used to retrieve spatial details lost during the process;
[0115] The decoder includes two branches that perform sliding window analysis to gradually restore the resolution of the feature map, perform detail reconstruction, and output image segmentation results.
[0116] Furthermore, the SM-UNet image segmentation model is built on the VMamba architecture, which has the following characteristics:
[0117] - Based on the State-Space Model (SSM), it processes sequence data through a scanning mechanism;
[0118] -Utilizes the properties of convolution to achieve a global receptive field, which is good at capturing long-range dependencies;
[0119] - A cross-scan strategy is used to process two-dimensional spatial relationships;
[0120] - It has the advantage of linear computational complexity.
[0121] However, SSM still has shortcomings in fine-grained localization and suffers from the drawback of information not being shared between windows. Therefore, this invention introduces the sliding window mechanism into the VMamba architecture, forming the SM-UNet image segmentation model of this invention. Combining the two, the SM-UNet image segmentation model of this invention possesses the following mechanisms:
[0122] 1. Sliding window mechanism for local enhancement of SSM:
[0123] Traditional state-space models (SSMs) construct global sequence dependencies through a scanning mechanism, but their continuous scanning characteristic faces two inherent challenges in vision tasks:
[0124] (1) Local detail weakening: The global scanning strategy of SSM (such as cross scanning) unfolds the two-dimensional space into a one-dimensional sequence, which may blur the local spatial relationship between pixels. For example, in edge detection tasks, the abrupt features of adjacent pixels may be diluted by the continuity of the scanning process.
[0125] (2) Information silos between windows: Although the global parameter sharing mechanism of the native SSM can capture long-range dependencies, it lacks explicit local inductive bias. When dealing with complex textures, similar patterns in different windows may not interact effectively, leading to duplicate calculations.
[0126] The sliding window mechanism of this invention addresses the above-mentioned shortcomings in the following ways:
[0127] (1) Local Feature Focusing: This method forces computation to be confined to a fixed window, preserving the integrity of the local structure through hard spatial partitioning. Experiments show that introducing a sliding window can reduce localization errors in keypoint detection tasks because it enhances local correlation modeling.
[0128] (2) Controllable cross-window interaction: Gradient permeation between adjacent windows is achieved through a shifting window strategy. The shifting operation causes the neurons at the window boundary to be assigned to different windows in the two calculations, forming an implicit information bridge. Mathematically, this is equivalent to adding a cross-window correlation constraint term to the loss function.
[0129] (3) Hierarchical feature distillation: Combined with hierarchical downsampling, the sliding window forms a pyramid-shaped feature abstraction. The shallow small window (4×4) captures details such as lines / corners, while the deep large window (14×14) integrates component-level semantics, which complements the single-scale scanning of SSM.
[0130] 2. SSM's global compensation mechanism for sliding windows:
[0131] The pure sliding window architecture has inherent limitations:
[0132] (1) Limited field of view: A single window can only capture local context. For tasks that require global reasoning, it may misjudge samples that are similar locally but contradictory globally.
[0133] (2) Long-range modeling is inefficient: it relies on layer-by-layer window shifting to pass information, and the theoretical receptive field expansion speed is slow, while SSM can achieve fast global coverage through state transfer.
[0134] The compensation of the SSM in this invention is reflected in global state memory and dynamic receptive field adjustment. Through cross-modal feature fusion, a direction-aware scanning strategy is designed, and the SSM can establish correlations between different scanning paths. For example, in medical image analysis, scanning paths along blood vessel orientation can enhance the global connectivity modeling of lesion regions, which dynamically complements the rigid division of the sliding window.
[0135] Example 2
[0136] This embodiment is a specific implementation of the SM-UNet image segmentation model in the image segmentation method based on the SM-UNet image segmentation model described in Embodiment 1.
[0137] The SM-UNet image segmentation model includes a convolutional module, an encoder, a skip connection module, and a decoder. For example... Figure 2As shown, the input 2D grayscale image has dimensions of H×W×1. It is first segmented into patches similar to those in ViT and VMamba by a convolutional module, and then converted into a 1D sequence. The initial Patch Embedding segmentation layer adjusts the feature dimension to an arbitrary size, denoted as C. These patch labels are then processed by multiple SwinVSS blocks and a Patch Merging layer to create hierarchical features. The Patch Merging layer handles downsampling and dimensionality increase, while the SwinVSS blocks focus on learning feature representations. The output resolutions of each stage of the encoder are H / 4×W / 4×C, H / 8×W / 8×2C, H / 16×W / 16×4C, and H / 32×W / 32×8C, respectively. The decoder consists of SwinVSS blocks and patch expansion layers following the encoder's style, capable of outputting the exact same feature size, thereby enhancing the spatial details lost during downsampling through skip connections. Two SwinVSS blocks are used in both the encoder and decoder.
[0138] Furthermore, the convolution module is used to perform convolutional segmentation processing on the image to be segmented, and output several feature map patches; the convolution module includes a convolutional layer and a Patch Embedding segmentation layer; the convolutional layer is used to convolve the image to be segmented into feature maps; the Patch Embedding segmentation layer is used to segment the feature maps and embed these small blocks into a high-dimensional space, and output several feature map patches of the same size.
[0139] Convolutional layer: Uses a convolution kernel to perform a convolution operation on the input image. The calculation formula is:
[0140]
[0141] Where I is the input image, W is the convolution kernel, Y is the output feature map, F is the convolution kernel size, and C is the input image. in This refers to the number of input channels. The convolution kernel slides across the input image, with stride S controlling the sliding speed and padding P used to maintain the output size. After the convolution operation, an activation function (such as ReLU) is typically applied to add non-linearity. The final output feature map serves as the input for the next layer, continuing with feature extraction or classification tasks.
[0142] The encoder includes two branches performing sliding window analysis to progressively extract and compress features from several feature map patches. The encoder also includes several SwinVSS blocks and a Patch Merging layer. The SwinVSS blocks have a self-attention mechanism to capture global and local features. The Patch Merging layer merges adjacent feature map patches, thereby reducing the resolution of the feature maps while increasing the feature dimension to capture higher-level features.
[0143] like Figure 3 As shown, the SwinVSS Block includes a first VSS Block and a second VSS Block connected in sequence.
[0144] The first VSS block includes a normalized (Layer Norm) layer and a fully connected (Linear) layer connected in sequence; wherein, a first branch and a second branch are provided between the normalized layer and the fully connected layer; the first branch includes a fully connected layer, a depth-wise convolutional layer, a W-2D selective scan layer, and a normalized (Layer Norm) layer connected in sequence; the second branch includes a fully connected layer;
[0145] The second VSS block includes a normalized (Layer Norm) layer and a fully connected (Linear) layer connected in sequence; wherein, a third branch and a fourth branch are provided between the normalized layer and the fully connected layer; the third branch includes a fully connected layer, a depth-wise convolutional layer, an SW-2D selective scan layer, and a normalized (Layer Norm) layer connected in sequence; the fourth branch includes a fully connected layer.
[0146] The W-2D selective scanning layer is used to divide the feature map patch that makes up the image to be segmented into M*M calculation windows.
[0147] The SW-2D selective scanning layer is used to slide the calculation window of the W-2D selective scanning layer by 1 / 2 window unit length in both the height and width dimensions.
[0148] The output expression of the W-2D selective scan layer in the first VSS block is:
[0149] z l =W-SS2D(LN(z) l-1 ))+z l-1 ,
[0150] Among them, z l For the output of the W-2D selective scan layer, z l-1SS2D() is the input for the W-2D selective scan layer, where W is the initial window selection, SS2D() is the 2D selective scan, and LN() is the normalization process.
[0151] The output expression of the SW-2D selective scan layer in the second VSS block is:
[0152] Z l+1 =SW-SS2D(LN(z) l ))+z l ,
[0153] Among them, z l+1 This is the output of the SW-2D selective scan layer, where SW is a sliding window processing layer.
[0154] Specifically, traditional state-space models (SSMs), as linear time-invariant system functions, will... Shot Through a hidden state Given As an evolutionary parameter, B, As a projection parameter of the state size, and with skip connections. This model can be represented as linear ordinary differential equations (ODEs), as shown in the formula:
[0155] h′(t)=Ah(t)+Bx(t),
[0156] y(t)=Ch(t)+Dx(t).
[0157] The discrete version of this linear model can be obtained by a zero-order preserves transformation, given a time-scale parameter.
[0158] h t =Ah k-2 +Bx k
[0159] y t =Ch k +Dx k
[0160] A = e ΔA
[0161] B = (e ΔA -I)A -1 B
[0162] C = C
[0163] in, Using a first-order Taylor series to approximate B, we get B = (e ΔA -I)A -1 B≈(ΔA)(ΔA)-1 ΔB = ΔB.
[0164] The depth-separable convolutional layer comprises 8 convolutional layers, with the number of channels in each layer being [16, 32, 64, 128, 128, 64, 32, 16].
[0165] Visual Mamba further introduces a cross-scan module (CSM) and then integrates convolutional operations into this module. In the VSS block, the input features first encounter a linear embedding layer and then branch into two paths. One branch passes through a depthwise separable convolution and SiLU activation, then enters the WSS2D module, and after post-layer normalization, merges with the other branch after SiLU activation. Unlike typical visual transformers, this SwinVSS block does not use positional embeddings but instead employs a streamlined structure that removes the MLP stage, thus enabling denser block stacking within the same depth budget.
[0166] In the standard Transformer architecture, each label needs to compute its relationship with all other labels, where the computational complexity is quadratic with the number of labels, making it unacceptable for many intensive prediction and high-resolution image tasks.
[0167] For efficient modeling, the SMUNet described in this embodiment uses window-based SS2D (WSS2D) and shift-window-based SS2D (SWSS2D). For example... Figure 4 The diagram shows a shift window method for computing self-attention in the proposed Swin Transformer architecture. In layer 1 (e.g.) Figure 4 As shown in a), a conventional window partitioning scheme is adopted, and visual Mamba computation is performed within each window. In the next layer (such as...), a conventional window partitioning scheme is used, and visual Mamba computation is performed within each window. Figure 4 As shown in b), the window partitioning shifts, creating a new window. The visual Mamba computation within the new window crosses the boundaries of the previous window in Layer 1, establishing a connection between them. The SS2D workflow is as follows: Figure 5 As shown, the input patch traverses along four different scan paths (cross-scan), with each sequence processed independently by a separate S6 module. The results are then merged to construct the final two-dimensional feature map (cross-merging).
[0168] In WSS2D, the input features are divided into non-overlapping windows, each containing M×M patches. WSS2D performs visual Mamba computation in only two directions within a local window. Assume z... l The output of the l-th layer WSS2D is calculated as follows:
[0169] z l =W-SS2D(LN(z)l-1 ))+z l-1 .
[0170] The problem with WMSA is the lack of effective information exchange between windows. To introduce cross-window interaction without adding extra computation, SWMSA is used. SWMSA's window configuration differs from previous WMSA layers; it uses an efficient batching method by cyclically shifting windows to the top left corner. After this shift, a batch window may consist of multiple non-adjacent sub-windows in the feature map, maintaining the same number of batch windows as in regular partitioning. When performing visual Mamba computation within local windows in both WMSA and SWMSA, the relative positional deviation is included in the similarity calculation. Through this shifted window partitioning mechanism, the output of the SWSS2D module can be written as:
[0171] z l+1 =SW-SS2D(LN(z) l ))+z l .
[0172] In the encoder, the C-dimensional tokenized input is processed for feature learning at reduced resolution through two consecutive SwinVSS blocks, preserving both dimensionality and resolution. The encoder uses triple patch merging as a downsampling process, dividing the input into quarter regions, concatenating them, and then normalizing the dimension each time through layernorm, halving the number of tokens while doubling the feature dimension.
[0173] The jump connection module is used to acquire spatial details lost during the process; the jump connection module includes two SwinVSS blocks.
[0174] Two SwinVSS blocks are used for the bottleneck section. Each level of the encoder and decoder employs skip connections to fuse multi-scale features with the upsampled output, enhancing spatial detail by merging shallow and deep layers. Subsequent linear layers maintain the dimensionality of this integrated feature set, ensuring consistency with the upsampled resolution.
[0175] The decoder includes two branches that perform sliding window analysis to gradually restore the resolution of the feature map, perform detail reconstruction, and output image segmentation results.
[0176] The decoder includes several Patch Expanding layers and a WinVSS Block.
[0177] The Patch Expanding layer is used to restore the resolution of the feature map and increase the number of feature map patches in order to perform more detailed feature extraction.
[0178] Similar to the encoder, the decoder uses two consecutive SwinVSS blocks for feature reconstruction, employing patch expansion layers instead of merging layers to amplify deep features. These layers halve the feature dimensions through the initial layer and then double the feature dimensions before reorganizing and reducing them to improve resolution, thus enhancing resolution.
[0179] In summary, this invention applies a sliding window to the Mamba-based UNet segmentation model. By segmenting the image into multiple local windows and performing calculations within these windows, the computational load is significantly reduced. By fully utilizing the cross-window interaction capabilities provided by the sliding window, global modeling capabilities are achieved to a certain extent while maintaining the efficiency of local computations.
[0180] Example 3
[0181] This embodiment is a practical application of the image segmentation method based on the SM-UNet image segmentation model described in Embodiment 2. The specific process is as follows:
[0182] 1. Dataset
[0183] This embodiment uses the publicly available ACDC magnetic resonance imaging (MRI) cardiac segmentation dataset from the MICCAI 2017 challenge. This dataset contains MRI scans of 100 patients and annotates various cardiac structures, such as the right ventricle and the endocardium and epicardium of the left ventricle. It covers a variety of pathological conditions and is categorized into five subgroups: normal, myocardial infarction, dilated cardiomyopathy, hypertrophic cardiomyopathy, and abnormal right ventricle, ensuring a broad distribution of feature information. Validation was performed on four regions of interest (ROIs) within the ACDC dataset.
[0184] 2. Implementation details
[0185] This example was performed on an Ubuntu 20.04 system using Python 3.8.8, PyTorch 1.10, and CUDA 11.3. The hardware setup included an Nvidia GeForce RTX 4080 GPU and an Intel Core i9-10900K CPU. For the ACDC dataset, the average runtime was approximately 5 hours, including data transfer, model training, and inference. The dataset is specifically designed for 2D image segmentation. The model was trained for 10,000 iterations with a batch size of 24. A stochastic gradient descent (SGD) optimizer was used with a learning rate of 0.01, momentum of 0.9, and weight decay of 0.0001. The network performance was evaluated on the validation set every 200 iterations, and model weights were only saved if a new best performance was achieved on the validation set.
[0186] 3. Evaluation Indicators
[0187] The evaluation of SWin-Mamba-UNet and baseline methods used a wide range of evaluation metrics. Similarity measures, preferably higher values, included: Dice, Intersection over Union (IoU), accuracy, precision, sensitivity, and specificity, indicated by upward arrows (↑), suggesting that higher values indicate better performance. Conversely, dissimilarity measures such as Hausdorff distance (HD) 95% and average surface distance (ASD), indicated by downward arrows (↓), were preferred at lower values, indicating a closer similarity between the predicted and ground truth segments.
[0188]
[0189] Where TP represents the number of true positives, TN represents the number of true negatives, FP represents the number of false positives, and FN represents the number of false negatives.
[0190]
[0191] Where a and b represent the sets of points on the predicted and true surfaces, respectively. d(a,b) represents the Euclidean distance between the two points.
[0192] 4. Results
[0193]
[0194] The experimental results are shown in the table above. According to the quantitative results, the SMUNet image segmentation model described in this invention is more likely to predict accurate segmentation masks.
[0195] Example 4
[0196] like Figure 6 As shown, an image segmentation device based on the SM-UNet image segmentation model includes at least one processor, a memory communicatively connected to the at least one processor, and at least one input / output interface communicatively connected to the at least one processor. The memory stores instructions executable by the at least one processor, which, when executed, enables the at least one processor to perform the image segmentation method based on the SM-UNet image segmentation model described in the foregoing embodiments. The input / output interface may include a display, keyboard, mouse, and USB interface for inputting and outputting data.
[0197] Those skilled in the art will understand that all or part of the steps of the above method embodiments can be implemented by hardware related to program instructions. The aforementioned program can be stored in a computer-readable storage medium. When the program is executed, it performs the steps of the above method embodiments. The aforementioned storage medium includes various media that can store program code, such as mobile storage devices, read-only memory (ROM), magnetic disks, or optical disks.
[0198] When the integrated units of this invention are implemented as software functional units and sold or used as independent products, they can also be stored in a computer-readable storage medium. Based on this understanding, the technical solutions of the embodiments of this invention, or the parts that contribute to the prior art, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the methods described in the various embodiments of this invention. The aforementioned storage medium includes various media capable of storing program code, such as mobile storage devices, ROMs, magnetic disks, or optical disks.
[0199] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the protection scope of the present invention.
Claims
1. An image segmentation method based on the SM-UNet image segmentation model, characterized in that, Includes the following steps: Obtain the image to be segmented; The image to be segmented is input into a pre-built SM-UNet image segmentation model; The SM-UNet image segmentation model outputs image segmentation results; The SM-UNet image segmentation model includes a convolutional module, an encoder, a skip connection module, and a decoder. The convolution module is used to perform convolutional segmentation processing on the image to be segmented and output several feature map patches. The encoder includes two branches for sliding window analysis, used to progressively extract and compress features from several feature map patches. The jump connection module is used to retrieve spatial details lost during the process; The decoder includes two branches for sliding window analysis, used to gradually restore the resolution of the feature map, perform detail reconstruction, and output image segmentation results; The encoder includes several SwinVSS blocks and a patch merging layer; The SwinVSS Block is used to capture global and local features; The Patch Merging layer is used to merge adjacent feature map patches; The SwinVSS Block includes a first VSS Block and a second VSS Block connected in sequence. The first VSS block includes a normalized layer and a fully connected layer connected in sequence; wherein, a first branch and a second branch are provided between the normalized layer and the fully connected layer; the first branch includes a fully connected layer, a depth-separable convolutional layer, a W-2D selective scan layer, and a normalized layer connected in sequence; the second branch includes a fully connected layer; The second VSS block includes a normalized layer and a fully connected layer connected in sequence; wherein, a third branch and a fourth branch are provided between the normalized layer and the fully connected layer; the third branch includes a fully connected layer, a depth-separable convolutional layer, an SW-2D selective scan layer, and a normalized layer connected in sequence; the fourth branch includes a fully connected layer; The W-2D selective scanning layer is used to divide the feature map patch that makes up the image to be segmented into M*M calculation windows. The SW-2D selective scanning layer is used to slide the calculation window of the W-2D selective scanning layer by 1 / 2 window unit length in both the height and width dimensions.
2. The image segmentation method based on the SM-UNet image segmentation model according to claim 1, characterized in that, The convolutional module includes convolutional layers and patch embedding segmentation layers; The convolutional layer is used to convolve the image to be segmented into a feature map; The Patch Embedding segmentation layer is used to segment the feature map into several feature map patches of the same size.
3. The image segmentation method based on the SM-UNet image segmentation model according to claim 2, characterized in that, The output expression of the W-2D selective scan layer in the first VSS block is: z l =W-SS2D(LN(z l-1 ))+z l-1 , Among them, z l For the output of the W-2D selective scan layer, z l-1 SS2D() is the input for the W-2D selective scan layer, where W is the initial window selection, SS2D() is the 2D selective scan, and LN() is the normalization process. The output expression of the SW-2D selective scan layer in the second VSS block is: WITH l+1 =SW-SS2D(LN(z l ))+z l , Among them, z l+1 This is the output of the SW-2D selective scan layer, where SW is a sliding window processing layer.
4. The image segmentation method based on the SM-UNet image segmentation model according to claim 3, characterized in that, The depth-separable convolutional layer comprises 8 convolutional layers, with the number of channels in each layer being [16, 32, 64, 128, 128, 64, 32, 16].
5. The image segmentation method based on the SM-UNet image segmentation model according to claim 3, characterized in that, The decoder includes several Patch Expanding layers and a WinVSS Block. The Patch Expanding layer is used to restore the resolution of the feature map and increase the number of feature map patches.
6. The image segmentation method based on the SM-UNet image segmentation model according to claim 3, characterized in that, The jump connection module includes two SwinVSS blocks.
7. An image segmentation device based on the SM-UNet image segmentation model, characterized in that, It includes at least one processor and a memory communicatively connected to the at least one processor; the memory stores instructions executable by the at least one processor to enable the at least one processor to perform the method of any one of claims 1 to 6.
Citation Information
Patent Citations
Jump connection method based on combination of progressive contraction strategy and shift window
CN117670904A
Road surface crack image segmentation method based on improved encoder-decoder structure
CN118736215A