Medical image segmentation method based on improved SwinUNet

Abstract

The invention relates to the technical field of medical image segmentation, in particular to a medical image segmentation method based on improved SwinUNet. According to the technical scheme, the method comprises the following steps that CT / MRI medical image data of an abdomen or heart area is collected and preprocessed, and preprocessing comprises image gray normalization, size resampling to 224 * 224 and enhancement processing; the preprocessed medical image is divided into image patch blocks, and initial feature representation is formed through linear embedding; through the improvement, the method is superior to a traditional convolutional neural network or a basic SwinUNet model in the aspects of segmentation precision, robustness, boundary processing, clinical applicability and the like, the problems of boundary discontinuity, poor small target recognition, insufficient cross-layer semantic alignment and the like existing in an existing method are effectively solved, a more reliable scheme is provided for automatic segmentation of medical images, and the method is suitable for large-scale popularization and application. The obvious clinical application value is realized.

Description

Technical Field

[0001] The present invention relates to the technical field of medical image segmentation, and in particular to a medical image segmentation method based on an improved SwinUNet. Background Art

[0002] With the advancement of medical imaging technology, the number and variety of medical images acquired clinically have increased rapidly. The ability to quickly and accurately segment anatomical structures in images through computer-assisted 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 handcrafted features and are difficult to adapt to complex clinical image data.

[0003] In recent years, deep learning methods, particularly convolutional neural networks (CNNs), have made significant progress in medical image segmentation. While architectures such as U-Net have performed well on multiple datasets, they are limited by their fixed convolution kernel structure and limited receptive field, making it difficult to effectively capture dependencies between long-range regions in an image.

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

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

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

[0007] Step 1: Collect CT / MRI medical imaging data of the abdomen or heart area and perform preprocessing;

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

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

[0010] Step 4: Introduce the Atrous Spatial Pyramid Pooling (ASPP) module at the end of the encoder to enhance the modeling capability of organs of different scales;

[0011] Step 5: The feature maps of the encoder multi-layer output are transmitted to the decoder through the skip connection, and the transformer in the corresponding skip path is inserted into the fusion module (TIF);

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

[0013] In step 7, a composite loss function is used during model training. The composite loss function includes Dice loss and boundary-enhanced cross entropy loss (BEL), and the SGD optimizer and cosine annealing strategy are used to dynamically adjust the learning rate.

[0014] Optionally, in step one, the preprocessing includes image grayscale normalization, resampling to 224×224, and enhancement processing, and 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 block in step 2 is 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 3 adopts a state space modeling method to replace the sliding window transformer (Swin Transformer), which effectively reduces the computational complexity and enhances the modeling capabilities of local and global features.

[0017] Optionally, the atrous spatial pyramid pooling module (ASPP) in step 4 is composed of convolution 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 vertical 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. The output feature map of the global average pooling branch is fused through concatenation and then compressed through a 1×1 convolution to obtain multi-scale context information, thereby better perceiving the size changes of the target structure and enhancing the network's generalization ability and segmentation accuracy.

[0023] Optionally, the transformer of the transformer insertion fusion module (TIF) 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 jump connection (sizes are (H / 2, W / 2, C) and (H, W, C)) and inputs them into an independent transformer layer to establish cross-resolution contextual dependencies, and then realizes weight reconstruction and semantic fusion through feature splicing and convolution operations, and finally outputs an enhanced feature representation (H, W, C).

[0024] Optionally, the encoder in step three is composed of a patch embedding module, a visual state space module (VSS) encoding layer, an image block merging module (Patch Merging) downsampling module and an atrous spatial pyramid pooling module (ASPP) stacked in sequence; the jump connection transfers feature maps between three scales of 1 / 4, 1 / 8 and 1 / 16.

[0025] Optionally, the decoder in step six is ​​composed of a Patch expansion module, a visual state space module (VSS) encoding layer, a feature fusion module and a final convolution output layer; each layer of the decoder includes a combination structure of an upsampling operation and a 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 the cross entropy loss, so that the model focuses on the segmentation performance of the boundary area of ​​the anatomical structure.

[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] Replacing the Swin Transformer block with state-space modeling enhances the model's ability to express long-range dependencies, solves the problem of broken local window attention mechanisms in traditional models, and makes feature expression more continuous.

[0030] The Atrous Spatial Pyramid Pooling (ASPP) module at the end of the encoder uses multiple atrous convolution paths with different atrous rates and a global average pooling path, along with cascade fusion and global context enhancement strategies, to improve the modeling capabilities of organs of different scales, especially small target recognition.

[0031] Through the Transformer Insertion Fusion Module (TIF) in the jump connection, this module aims to solve the problems of insufficient semantic alignment, information redundancy and missing details in the traditional feature splicing method, 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 area of ​​anatomical structures and solves the problem of insufficient continuity of organ boundaries.

[0033] This invention is suitable for 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 enhancement strategies improve the model's adaptability to different imaging conditions and have high potential for clinical application.

[0034] Through the above-mentioned improvements, the present invention outperforms traditional convolutional neural networks or the basic SwinUNet model in terms of segmentation accuracy, robustness, boundary processing, and clinical applicability. It effectively solves the problems of boundary discontinuity, poor small target recognition, and insufficient cross-layer semantic alignment in existing methods, providing a more reliable solution for automated segmentation of medical images and having significant clinical application value. BRIEF DESCRIPTION OF THE DRAWINGS

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

[0036] Figure 2 To improve the SwinUNet model network structure diagram according to the embodiment of the present invention.

[0037] Figure 3 2 is a structural diagram of the spatial hollow pyramid ASPP module according to an embodiment of the present invention.

[0038] Figure 4 Detailed structural diagram of a visual state space (VSS) module according to an embodiment of the present invention.

[0039] Figure 5 FIG. 4 is a structural diagram of a transformer insertion fusion module (TIF) according to an embodiment of the present invention. DETAILED DESCRIPTION

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

[0041] Example

[0042] like Figure 1 As shown, the present invention provides an improved SwinUNet medical image segmentation method, and each step of the method is described in detail below.

[0043] 1. Collect and preprocess medical imaging data

[0044] This paper targets the automated segmentation of medical images of multiple abdominal organs and the cardiac region, suitable for CT and MRI data in DICOM or NIfTI formats. After data acquisition, the images are first grayscale normalized, normalizing pixel intensities to the [0, 1] range to ensure numerical stability during model training. Second, the images are resampled to a uniform size of 224×224 pixels to accommodate the network input size requirements. Furthermore, to improve the model's robustness and generalization capabilities, the training data is augmented using the following data augmentation techniques:

[0045] 1) Random rotation of ±15° to simulate slight deviation of the patient's scanning angle;

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

[0047] 3) Brightness variation of ±20%, enhancing the stability of the model under different imaging conditions;

[0048] 4) Mirror flipping and random occlusion improve the adaptability to symmetrical structures and missing information.

[0049] 2. Image segmentation and embedding initial features

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

[0051] 3. Encoder structure and VSS module design

[0052] As a key component of feature extraction, the encoder's overall architecture draws inspiration from the multi-scale design of U-Net. The encoder consists of multiple stacked stages, each of which includes a patch merging module and a variable-scale semantic module (VSS module). The patch merging module merges adjacent image patches, halving the spatial resolution and doubling the number of channels, 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 architecture is based on state-space modeling theory and employs a set of learnable state variables and a recursive update mechanism to model dependencies between long-range pixels in an image. Unlike traditional Transformer architectures, 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 layer are determined by the feature changes of the current stage. During the feature update process, VSS selectively fuses the historical state with the current input through a gating mechanism to obtain a more robust global feature expression.

[0055] 4. ASPP module and multi-scale enhancement

[0056] like Figure 3 As shown in the figure, the present invention introduces a dilated spatial pyramid pooling module (ASPP) at the end of the encoder to improve the model's ability to perceive objects 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, respectively, and can perceive contextual information in neighborhoods 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 limited ability to model long-range dependencies. The feature maps output by each branch are concatenated in the channel dimension and fused through a 1×1 convolution to form a feature map with a uniform number of channels. This serves as the encoder's final output and enters the Bottleneck layer.

[0058] 5. Skip Connection and TIF Insertion Fusion Mechanism

[0059] Traditional U-Net uses feature splicing to connect the encoder and decoder, which is prone to semantic inconsistency. This paper introduces the TIF module at the skip connection as a fuser on the skip connection path, effectively alleviating the semantic gap between shallow and deep features.

[0060] The TIF module structure is as shown in the attached 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, first unifies the channels through a (1, 1, c) convolution. After processing the sequenced features through independent Transformer layers, the output is restored to its spatial dimensions and concatenated into a (h, w, c2) tensor. Finally, compressed convolution is performed to generate the fused features (h, w, c). The core design uses Transformers to achieve cross-resolution feature interaction while maintaining spatial compatibility between the output and input.

[0061] 6. 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, gradually restoring spatial resolution. In each decoding stage, upsampling by a factor of 2 is first performed through nearest neighbor interpolation or transposed convolution. Features from skip connections are then integrated and fused via a TIF module. Finally, a VSS module further refines 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 (such as the number of organs + background) and outputs the final segmentation prediction through the softmax layer.

[0064] 7. Loss Function and Training Strategy

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

[0066] 1) Dice loss: measures the regional overlap between predictions and labels, suitable for dealing with sample imbalance problems;

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

[0068] The training optimizer used SGD, with momentum set to 0.9 and weight decay set to 0.0001. The initial learning rate was 0.05, and the learning rate was dynamically decreased using a cosine annealing strategy. The training batch size was set to 24, with a total of 150 epochs. Multi-GPU parallel acceleration was supported.

[0069] Experimental setup and evaluation metrics

[0070] A 5-fold cross-validation method is used for training and testing. The evaluation indicators include:

[0071] 1) Dice coefficient (DSC): measures the degree of overlap between the predicted area and the true 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 and consists of an encoder, a decoder and jump connections.

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

[0076] The decoder input is further processed by the Bottleneck layer, using VSS and ASPP to enhance multi-scale feature extraction and prepare the segmentation output. The output is then further processed by Patch Merging and VSS, gradually restoring the resolution of the feature maps. Each decoder layer incorporates upsampling and feature fusion to help recover detailed information. Furthermore, each layer incorporates VSS, which helps preserve features from the encoder, thereby improving the quality of feature map segmentation.

[0077] In this embodiment, skip connections enable the appropriate fusion of low-level, high-resolution features with high-level, low-resolution features, providing high-quality feature information for the final segmentation result. The TIF module in skip connections further optimizes the feature fusion process, particularly when fusing high-resolution and low-resolution features. It adaptively adjusts weights to strengthen key features, enhancing more important features and ultimately improving the final segmentation result.

[0078] As a medical image segmentation method based on the improved SwinUNet, the present invention has many significant benefits in the field of medical image processing. By introducing the visual state space module (VSS) and replacing the Swin Transformer block with a state space modeling method, the model's ability to express long-range dependencies is enhanced, making feature expression more continuous, effectively improving the problem of broken local window attention mechanism in traditional models, and improving the integrity and accuracy of feature extraction.

[0079] Furthermore, the atrous spatial pyramid pooling module (ASPP) at the end of the encoder contains multiple atrous convolution paths with different atrous rates (1, 3, 5, 7) and a global average pooling path, which enhances the modeling ability of organs of different scales, especially improving the recognition effect of small target organs, and solving the problem of insufficient recognition ability of small organs in existing methods.

[0080] In addition, a fusion module (TIF) is inserted into the jump connection to solve the problems of insufficient semantic alignment, information redundancy and missing details in the traditional feature splicing method, and to enhance the semantic interaction and fine-grained integration capabilities between features at different levels.

[0081] The present invention optimizes the boundary segmentation effect: in the composite loss function used in model training, the boundary-enhanced cross-entropy loss (BEL) 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 boundary area of ​​the anatomical structure, 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 clinical requirements for boundary accuracy.

[0082] It is worth mentioning that the adaptability to clinical applications is enhanced: this method 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 cope with 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, and enhance the model's generalization and robustness, enabling it to function stably in diverse clinical data scenarios.

[0084] Through the above improvements, this method outperforms traditional convolutional neural networks (CNN) 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 clinicians, assist doctors in improving diagnostic efficiency and accuracy, reduce the workload of manual outlining, and has high clinical application potential and promotion value.

[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 inspirations of the above embodiments, those skilled in the art may make various alternative improvements and combinations to the above specific embodiments.

Claims

1. A medical image segmentation method based on improved SwinUNet, characterized in that: The following steps are involved: Step 1: Collect CT / MRI medical imaging data of the abdomen or heart area and perform preprocessing; Step 2: Divide the preprocessed medical image into image patches and form the initial feature representation through linear embedding; Step 3: Using an encoder that introduces a visual state space module to extract features from the image patch and construct multi-scale context information; Step 4: A dilated spatial pyramid pooling module is introduced at the end of the encoder in the model to enhance the modeling capability of organs at different scales. Step 5: The feature maps of the encoder multi-layer output in the model are transmitted to the decoder through the skip connection, and the transformer in the corresponding skip path is inserted into the fusion module; Step 6: The model decoder uses the visual state space module and upsampling operations to restore the spatial resolution layer by layer and output the final segmented image; Step 7: A composite loss function is used during model training. The composite loss function includes Dice loss and boundary-enhanced cross entropy loss, and the SGD optimizer and cosine annealing strategy are used to dynamically adjust the learning rate.

2. The medical image segmentation method based on improved SwinUNet according to claim 1, characterized in that: In step one, the preprocessing includes image grayscale normalization, resampling to 224×224, and enhancement processing, and 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 improved SwinUNet according to claim 1, characterized in that: The image patch block in step 2 is 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 improved SwinUNet according to claim 1, characterized in that: The visual state space module in step 3 adopts a state space modeling approach to replace the sliding window transformer.

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

6. The medical image segmentation method based on improved SwinUNet according to claim 1, characterized in that: The transformer inserted into the fusion module in step 5 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 jump connection 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, to establish cross-resolution contextual dependencies, and then achieve weight reconstruction and semantic fusion through feature splicing and convolution operations, and finally output an enhanced feature representation (H, W, C).

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

8. The medical image segmentation method based on improved SwinUNet according to claim 1, characterized in that: The decoder in step six is ​​composed of a patch expansion module, a visual state space module encoding layer, a feature fusion module and a final convolution output layer; each layer of the decoder includes a combination structure of an upsampling operation and a TIF module.

9. The medical image segmentation method based on improved SwinUNet according to claim 1, characterized in that: The boundary enhancement cross entropy loss in step seven 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 boundary area of ​​the anatomical structure.

10. The medical image segmentation method based on improved SwinUNet according to claim 1, characterized in that: In step 7, 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