A Medical Image Segmentation Method Based on Improved SwinUNet

Abstract

This invention relates to the field of medical image segmentation technology, and more particularly to a medical image segmentation method based on an improved SwinUNet. The technical solution includes the following steps: acquiring CT / MRI medical image data of the abdominal or cardiac region, and performing preprocessing, including image grayscale normalization, resampling to 224×224, and enhancement processing; dividing the preprocessed medical image into image patch blocks, and forming initial feature representations through linear embedding. Through the above improvements, this invention outperforms traditional convolutional neural networks or basic SwinUNet models in terms of segmentation accuracy, robustness, boundary handling, and clinical applicability. It effectively solves problems such as boundary discontinuities, poor small target recognition, and insufficient cross-layer semantic alignment in existing methods, providing a more reliable solution for automated medical image segmentation and possessing significant clinical application value.

Description

Technical Field

[0001] This invention relates to the field of medical image segmentation technology, and in particular to a medical image segmentation method based on an improved SwinUNet. Background Technology

[0002] With the development of medical imaging technology, the number and types of medical images acquired in clinical settings have increased rapidly. How to quickly and accurately segment anatomical structures in images using computer-aided methods has become a core task in current medical image analysis research. Traditional image segmentation methods, such as thresholding, region growing, and graph-based segmentation algorithms, typically rely on manual features and are ill-suited to complex clinical image data.

[0003] In recent years, deep learning methods, especially convolutional neural networks (CNNs), have made significant progress in medical image segmentation tasks. Architectures represented by U-Net have performed well on multiple datasets, but due to their fixed convolutional kernel structure and limited receptive field, they are unable to effectively capture the dependencies between long-range regions in images.

[0004] To address this, researchers introduced the Transformer structure to enhance the model's global modeling capabilities. SwinUNet, a segmentation framework that embeds the Swin Transformer into a U-Net structure, effectively integrates local and global information. However, SwinUNet still suffers from issues such as broken local window attention mechanisms, insufficient semantic information alignment, and poor small target modeling capabilities. Therefore, we propose a medical image segmentation method based on an improved SwinUNet. Summary of the Invention

[0005] The purpose of this invention is to address the problems of broken local window attention mechanism, insufficient semantic information alignment, and poor small target modeling ability in the background art of SwinUNet, and to propose a medical image segmentation method based on an improved SwinUNet.

[0006] The technical solution of this invention: a medical image segmentation method based on an improved SwinUNet, comprising the following steps:

[0007] Step 1: Acquire CT / MRI medical imaging data of the abdominal or cardiac region and perform preprocessing;

[0008] Step 2: Divide the preprocessed medical image into image patch blocks and form initial feature representations through linear embedding;

[0009] Step 3: Use an encoder with a visual state space module (VSS) to extract features from the image patch and construct multi-scale context information;

[0010] Step 4: A void space pyramid pooling module (ASPP) is introduced at the end of the encoder to enhance the modeling capability of organs at different scales.

[0011] Step 5: The feature map output by the encoder's multi-layer output is transmitted to the decoder through a skip connection, and the transformer in the corresponding skip path is inserted into the fusion module (TIF).

[0012] Step six: The decoder uses the VSS module and upsampling operation to restore the spatial resolution layer by layer and outputs the final segmented image;

[0013] Step 7: During model training, a composite loss function is used, which includes Dice loss and boundary augmentation cross-entropy loss (BEL), and the learning rate is dynamically adjusted using an SGD optimizer and a cosine annealing strategy.

[0014] Optionally, in step one, the preprocessing includes image grayscale normalization, size resampling to 224×224, and enhancement processing. The enhancement processing includes at least one of random rotation, brightness adjustment, elastic deformation, mirror flipping, and random occlusion. The angle range of the random rotation is ±15°, and the range of the brightness adjustment is ±20%.

[0015] Optionally, the image patch blocks in step two are 4×4 in size, and each patch block is converted into a vector representation of uniform length through linear embedding.

[0016] Optionally, the Visual State Space Module (VSS) in step three uses a state space modeling approach instead of a sliding window transformer, which effectively reduces computational complexity and enhances the ability to model local and global features.

[0017] Optionally, the Spatial Pyramid Pooling (ASPP) module in step four is composed of convolutional kernels of multiple scales in parallel, including:

[0018] A standard 1×1 convolution;

[0019] A standard 3×3 convolution;

[0020] A 3×4 convolution with a dilation rate of 6 is used to stretch the longitudinal receptive field of the feature;

[0021] A 3×3 dilated convolution with a dilation rate of 12;

[0022] A global average pooling (Image Pooling) branch is used. The output feature maps of the global average pooling branch are fused through concatenation and then compressed through a 1×1 convolution to obtain multi-scale contextual information, thereby better perceiving changes in the size of the target structure and enhancing the network's generalization ability and segmentation accuracy.

[0023] Optionally, the transformer in the Transformer Insertion Fusion (TIF) module in step five is a neural network structure that uses a self-attention mechanism to model sequence relationships. The module serializes the multi-scale feature maps (with sizes of (H / 2, W / 2, C) and (H, W, C) respectively) in the skip connections and inputs them into independent transformer layers to establish cross-resolution contextual dependencies. Then, weight reconstruction and semantic fusion are achieved through feature concatenation and convolution operations, and finally, the enhanced feature representation (H, W, C) is output.

[0024] Optionally, the encoder in step three is composed of a Patch embedding module, a Visual State Space (VSS) encoding layer, a Patch Merging downsampling module, and an Aperture Spatial Pyramid Pooling (ASPP) module stacked sequentially; the skip connections transfer feature maps between three scales: 1 / 4, 1 / 8, and 1 / 16.

[0025] Optionally, the decoder in step six consists of a Patch extension module, a Visual State Space (VSS) encoding layer, a feature fusion module, and a final convolutional output layer; each layer of the decoder includes a combination structure of upsampling operation and TIF module.

[0026] Optionally, the Boundary Enhanced Cross-Entropy Loss (BEL) in step seven introduces a boundary weighting mechanism based on spatial distance on the basis of cross-entropy loss, so that the model focuses on the segmentation performance of the anatomical boundary region.

[0027] Optionally, in step seven, the momentum of the SGD optimizer is set to 0.9, and the weight decay is set to 0.0001; the initial learning rate of the cosine annealing strategy is 0.05.

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

[0029] By replacing the Swing Transformer block with state-space modeling, the model's ability to express long-range dependencies is enhanced, the problem of broken local window attention mechanisms in traditional models is solved, and feature representation is made more continuous.

[0030] The Spatial Pyramid Pooling Module (ASPP) at the end of the encoder enhances the modeling ability of organs at different scales, especially improving the recognition effect of small targets, through multiple dilated convolutional paths with different dilation rates and global average pooling paths, and through cascaded fusion and global context enhancement strategies.

[0031] By using the Transformer Insertion Fusion Module (TIF) in the skip connection, this module aims to solve the problems of insufficient semantic alignment, information redundancy and lack of details in traditional feature splicing methods, and enhance the semantic interaction and fine-grained integration capabilities between features at different levels.

[0032] The composite loss function used in model training includes Boundary Enhanced Cross-Entropy Loss (BEL). By introducing a boundary weighting mechanism based on spatial distance, the model focuses on the boundary regions of anatomical structures, thus solving the problem of insufficient continuity of organ boundaries.

[0033] This invention is applicable to the automated segmentation of medical images such as CT and MRI, and performs particularly well in complex areas such as brain tumors, cardiac structures, and multiple abdominal organs. The preprocessing and data augmentation strategies enhance the model's adaptability to different imaging conditions and have high potential for clinical application.

[0034] Through the above improvements, this invention outperforms traditional convolutional neural networks or basic SwinUNet models in terms of segmentation accuracy, robustness, boundary handling, and clinical applicability. It effectively solves the problems of boundary discontinuity, poor recognition of small targets, and insufficient cross-layer semantic alignment in existing methods, providing a more reliable solution for automated medical image segmentation and having significant clinical application value. Attached Figure Description

[0035] Figure 1 This is a flowchart of an improved SwinUNet medical image segmentation method according to an embodiment of the present invention.

[0036] Figure 2 The network structure diagram of the SwinUNet model is improved according to an embodiment of the present invention.

[0037] Figure 3 This is a structural diagram of the Space Void Pyramid (ASPP) module according to an embodiment of the present invention.

[0038] Figure 4 This is a detailed structural diagram of the Visual State Space (VSS) module according to an embodiment of the present invention.

[0039] Figure 5 This is a structural diagram of a transformer insertion fusion module (TIF) according to an embodiment of the present invention. Detailed Implementation

[0040] The technical solution of the present invention will now be clearly and completely described with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments.

[0041] Example

[0042] like Figure 1 As shown, this invention provides an improved SwinUNet medical image segmentation method, and the steps of the method are described in detail below.

[0043] I. Acquiring and Preprocessing Medical Imaging Data

[0044] This invention targets the automatic segmentation of medical images of multiple organs in the abdomen and the heart region, and is applicable to CT and MRI data in DICOM or NIfTI format. After data acquisition, the images are first normalized to grayscale, standardizing pixel intensity to the [0,1] range to ensure numerical stability during model training. Secondly, the images are resampled to a size of 224×224 pixels to fit the input size requirements of the network structure. Furthermore, to improve the robustness and generalization ability of the model, the following data augmentation techniques are used to expand the training data:

[0045] 1) Randomly rotate ±15° to simulate slight shifts in the patient's scanning angle;

[0046] 2) Elastic deformation, simulating the morphological changes of tissues and organs caused by respiratory movements;

[0047] 3) Brightness variation ±20% enhances model stability under different imaging conditions;

[0048] 4) Mirroring and random occlusion improve the ability to adapt to symmetrical structures and missing information.

[0049] II. Image segmentation and embedding of initial features

[0050] The preprocessed 2D image is divided into 4×4 image patches, each patch being 56×56 pixels in size. Before being input into the network, each patch undergoes dimensionality mapping via a linear embedding module, converting it into a vector representation of uniform length. This process not only preserves local structural information but also facilitates position-aware modeling by subsequent modules (such as VSS).

[0051] III. Encoder Structure and VSS Module Design

[0052] As a key component of feature extraction, the encoder's overall structure draws inspiration from the multi-scale design of U-Net. The encoder consists of multiple stacked stages, each containing a patch merging module and a variable-scale semantic module (VSS module). The patch merging module achieves halving of spatial resolution and doubling of the number of channels by merging adjacent image patches, thereby extracting higher-level abstract features.

[0053] like Figure 4 As shown, the VSS (Visual State Space) module is one of the core innovations of this invention. Its structure is based on state-space modeling theory, employing a set of learnable state variables and a recursive update mechanism to display the dependencies between distant pixels in the modeled image. Unlike the traditional Transformer structure, VSS avoids pairwise self-attention calculations for all pixel pairs, significantly reducing computational complexity.

[0054] The input channel dimension and spatial size of each VSS module are determined according to the feature changes of the current stage. During feature update, VSS selectively fuses historical states with the current input through a gating mechanism to obtain a more robust global feature representation.

[0055] IV. ASPP Module and Multi-Scale Enhancement

[0056] like Figure 3 As shown, this invention introduces a Spatial Pyramid Pooling (ASPP) module at the end of the encoder to enhance the model's ability to perceive targets of different scales (especially small organs). The ASPP module consists of four parallel dilated convolution branches with dilation rates of 1, 3, 5, and 7, and can perceive contextual information of a neighborhood of 1 to 7 pixels.

[0057] In addition, a global average pooling path is added to capture the global contextual features of the entire image, thereby compensating for the convolutional network's insufficient ability to model long-distance dependencies. The feature maps output by each branch are concatenated along the channel dimension and then fused into a feature map with a uniform number of channels through a 1×1 convolution, which serves as the final output of the encoder and enters the Bottleneck layer.

[0058] V. Skip Connections and TIF Insertion Fusion Mechanism

[0059] Traditional U-Net connects the encoder and decoder using feature concatenation, which can easily lead to semantic inconsistencies. This invention introduces a TIF module at the skip connection point as a fusion agent on the skip connection path, effectively mitigating the semantic gap between shallow and deep features.

[0060] The structure of the TIF module is shown in the attached figure. Figure 5As shown, the dual-branch Transformer fusion module (TIF) takes high- and low-resolution features ((h / 2, w / 2, c) and (h, w, c)) as input. It first unifies the channels through a (1, 1, c) convolution, then processes the serialized features through independent Transformer layers. The output is then restored to its spatial dimension and concatenated into a (h, w, c2) tensor. Finally, a compressed convolution generates the fused feature (h, w, c). The core design achieves cross-resolution feature interaction through Transformer while maintaining spatial compatibility between the output and input.

[0061] VI. Decoder Design and Layer-by-Layer Reconstruction

[0062] The decoder maintains a symmetrical structure with the encoder, and its task is to restore high-level abstract features to the original image size and gradually restore spatial resolution. In each decoding stage, the features are first upsampled by a factor of 2 through nearest neighbor interpolation or transposed convolution; then, the features from skip connections are incorporated and fused via a TIF module; finally, a VSS module further optimizes the feature representation.

[0063] The last layer of the decoder uses a 1×1 convolution to map the feature map to the number of categories (e.g., number of organs + background), and outputs the final segmentation prediction through a softmax layer.

[0064] VII. Loss Function and Training Strategy

[0065] The loss function used in this invention is a composite form, including:

[0066] 1) Dice loss: measures the degree of regional overlap between the prediction and the label, and is suitable for handling imbalanced sample problems;

[0067] 2) Boundary Enhanced Cross-Entropy Loss (BEL): Based on the regular cross-entropy loss, a boundary weighting function based on distance transformation is introduced to enhance the model's ability to focus on edge pixels and improve the segmentation accuracy of boundary details.

[0068] The training optimizer uses SGD with a momentum of 0.9 and a weight decay parameter of 0.0001. The initial learning rate is 0.05, dynamically decreasing using cosine annealing. The batch size is set to 24, with a total of 150 training epochs, and multi-GPU parallel acceleration is supported.

[0069] Experimental setup and evaluation indicators

[0070] Training and testing were performed using a 5-fold cross-validation method, and the evaluation metrics included:

[0071] 1) Dice coefficient (DSC): measures the degree of overlap between the predicted area and the actual area;

[0072] 2) Hausdorff distance (HD): Evaluates the maximum error distance between segmentation boundaries;

[0073] 3) IoU (Intersection over Union): Evaluates the accuracy of segmentation.

[0074] like Figure 2 As shown, it is worth noting that the overall architecture follows a U-shaped symmetrical structure, consisting of an encoder, a decoder, and jump connections.

[0075] The encoder divides the input image into 4×4 patch blocks, each block being 1 / 4 the size of the original image. Each patch first undergoes linear embedding, then is input into the VSS for feature extraction. This module typically includes multiple attention mechanisms and transformation layers to enhance the expressive power of local and global features. This process is repeated twice to more effectively extract features at different scales. As the network depth increases, the resolution of the feature maps gradually decreases. After each patch merging operation, the size of the feature map is halved, but the number of channels doubles, improving feature expressive power and helping the model extract higher-level semantic information. To ensure effective feature transfer and avoid information loss, skip connections are introduced between different layers. Each skip connection transfers feature sizes of 1 / 4, 1 / 8, and 1 / 16, depending on its layer.

[0076] Furthermore, the decoder input is first processed through the Bottleneck layer, where VSS and ASPP are used to further enhance multi-scale feature extraction and prepare for the segmentation output. After the output is processed, Patch Merging and VSS are used to further restore the resolution of the feature maps level by level. Each layer of the decoder includes upsampling and feature fusion processes, which help recover detailed information. Additionally, each layer also includes VSS, which helps preserve features from the encoder, thereby improving the segmentation quality of the feature maps.

[0077] In this embodiment, skip connections achieve appropriate fusion of low-level high-resolution features and high-level low-resolution features, providing high-quality feature information for the final segmentation result. The TIF module further optimizes the feature fusion process in skip connections, especially when fusing high-resolution and low-resolution features. By adaptively adjusting weights to strengthen key features, more important features are enhanced, thereby improving the final segmentation effect.

[0078] This invention, as a medical image segmentation method based on an improved SwinUNet, has many significant benefits in the field of medical image processing. By introducing a visual state space module (VSS) and replacing the Swin Transformer block with a state space modeling approach, it enhances the model's ability to express long-range dependencies, makes feature representation more continuous, effectively improves the problem of broken local window attention mechanisms in traditional models, and improves the completeness and accuracy of feature extraction.

[0079] Furthermore, the Spatial Pyramid Pooling Module (ASPP) at the end of the encoder contains multiple dilated convolutional paths with different dilation rates (1, 3, 5, 7) and a global average pooling path, which enhances the ability to model organs of different scales, especially improving the recognition effect of small target organs, and solving the problem of insufficient recognition ability of existing methods for small organs.

[0080] In addition, the insertion of a fusion module (TIF) in skip connections aims to solve the problems of insufficient semantic alignment, information redundancy and lack of details in traditional feature splicing methods, and enhance the semantic interaction and fine-grained integration capabilities between features at different levels.

[0081] This invention optimizes boundary segmentation performance: In the composite loss function used in model training, the Boundary Enhancement Cross-Entropy Loss (BEL) introduces a boundary weighting mechanism based on spatial distance on the basis of cross-entropy loss, so that the model focuses on the segmentation performance of the anatomical structure boundary region, effectively improving the problem of insufficient continuity of organ boundaries, improving the accuracy of image edge segmentation, and making the segmentation results more in line with the clinical requirements for boundary accuracy.

[0082] It is worth noting that this method enhances its adaptability to clinical applications: it is suitable for automated segmentation tasks of various medical images such as CT and MRI, and performs particularly well in the automatic delineation of complex areas such as brain tumors, cardiac structures and multiple abdominal organs, and can meet the segmentation challenges of different anatomical structures.

[0083] Image grayscale normalization, size resampling, and various enhancement processes (such as random rotation and brightness adjustment) in the preprocessing stage improve the model's adaptability to different imaging conditions, scanning angles, and tissue variations, enhancing the model's generalization and robustness, and enabling it to function stably in diverse clinical data scenarios.

[0084] With the above improvements, this method outperforms traditional convolutional neural networks (CNNs) or basic SwinUNet models in terms of organ boundary continuity, small target recognition, and cross-anatomical structure segmentation. It can provide more accurate and reliable automatic segmentation results for clinical use, help doctors improve diagnostic efficiency and accuracy, reduce the workload of manual delineation, and has high potential for clinical application and promotion.

[0085] The above specific embodiments are merely several optional embodiments of the present invention. Based on the technical solutions of the present invention and the relevant teachings of the above embodiments, those skilled in the art can make various alternative improvements and combinations to the above specific embodiments.

Claims

1. A medical image segmentation method based on an improved SwinUNet, characterized in that, Includes the following steps: Step 1: Acquire CT / MRI medical imaging data of the abdominal or cardiac region and perform preprocessing; Step 2: Divide the preprocessed medical image into image patch blocks and form initial feature representations through linear embedding; Step 3: The image patch is feature extracted using an encoder that incorporates a visual state space module to construct multi-scale context information. The visual state space module uses a state space modeling approach instead of a sliding window transformer. Step 4: Introduce a hollow space pyramid pooling module at the end of the encoder in the model to enhance the modeling ability of organs at different scales. Step 5: The feature maps of the encoder's multi-layer outputs in the model are transmitted to the decoder through skip connections, and the transformers in the corresponding skip paths are inserted into the fusion module. Step six: The model decoder uses a visual state space module and upsampling operations to restore the spatial resolution layer by layer, and outputs the final segmented image; Step 7: During model training, a composite loss function is used, which includes Dice loss and boundary-enhanced cross-entropy loss. The learning rate is dynamically adjusted using an SGD optimizer and a cosine annealing strategy. The boundary-enhanced cross-entropy loss introduces a boundary weighting mechanism based on spatial distance on the basis of the cross-entropy loss, so that the model focuses on the segmentation performance of the anatomical structure boundary region.

2. The medical image segmentation method based on the improved SwinUNet according to claim 1, characterized in that, In step one, the preprocessing includes image grayscale normalization, size resampling to 224×224, and enhancement processing. The enhancement processing includes at least one of random rotation, brightness adjustment, elastic deformation, mirror flipping, and random occlusion. The angle range of the random rotation is ±15°, and the range of the brightness adjustment is ±20%.

3. The medical image segmentation method based on the improved SwinUNet according to claim 1, characterized in that, The image patch blocks in step two are 4×4 in size, and each patch block is converted into a vector representation of uniform length through linear embedding.

4. The medical image segmentation method based on the improved SwinUNet according to claim 1, characterized in that, The hollow spatial pyramid pooling module in step four consists of parallel convolutional kernels of multiple scales, including: A standard 1×1 convolution; Standard 3×3 convolution; A 3×4 convolution with a dilation rate of 6 is used to stretch the longitudinal receptive field of the feature; 3×3 dilated convolution with a dilation rate of 12; The output feature map of the global average pooling branch is cascaded and then compressed by 1×1 convolution to obtain multi-scale contextual information.

5. A medical image segmentation method based on an improved SwinUNet according to claim 1, characterized in that, The transformer in the transformer insertion fusion module in step five is a neural network structure that uses a self-attention mechanism to model sequence relationships. The module serializes the multi-scale feature maps in the skip connections and inputs them into an independent transformer layer. The sizes of the multi-scale feature maps are (H / 2, W / 2, C) and (H, W, C), respectively. It establishes cross-resolution contextual dependencies, and then performs weight reconstruction and semantic fusion through feature concatenation and convolution operations, finally outputting an enhanced feature representation (H, W, C).

6. The medical image segmentation method based on the improved SwinUNet according to claim 1, characterized in that, The encoder in step three is composed of a Patch embedding module, a visual state space module encoding layer, an image patch merging module downsampling module, and an enhanced void space pyramid pooling module stacked sequentially; the skip connection transmits feature maps between three scales: 1 / 4, 1 / 8, and 1 / 16.

7. A medical image segmentation method based on an improved SwinUNet according to claim 1, characterized in that, The decoder in step six consists of a patch extension module, a visual state space module encoding layer, a feature fusion module, and a final convolutional output layer; each layer of the decoder contains a combination structure of upsampling operation and TIF module.

8. A medical image segmentation method based on an improved SwinUNet according to claim 1, characterized in that, In step seven, the momentum of the SGD optimizer is set to 0.9, the weight decay parameter is set to 0.0001, and the initial learning rate of the cosine annealing strategy is 0.05.

Images