Medical image segmentation method based on GM-TransUnet

Abstract

The invention discloses a medical image segmentation method based on GM-TransUnet. The medical image segmentation method mainly aims at solving the problems of leap connection semantic gaps and upsampling local information loss in an existing U-shaped network. According to the method, a TransUnet model is improved, and convolution local attention (CLATtn) and global-local attention (GLATtn) are added, so that the feature quality of jump connection transmission is enhanced, and the semantic gap problem of jump connection is effectively relieved. Besides, a multi-scale fusion attention module (MFAttn) is further added, and local information loss in the up-sampling process is remarkably reduced through a multi-scale feature fusion mechanism. The above modules cooperate with each other, so that the network is effectively improved in the aspects of boundary positioning and overall segmentation precision. The method is suitable for a clinical medical image segmentation scene, and has good practical application value and popularization prospect.

Description

Technical Field

[0001] This invention relates to the field of medical image segmentation, and more particularly to a medical image segmentation method based on GM-TransUnet. Background Technology

[0002] Medical image segmentation, a crucial task in computer vision and medical image analysis, aims to accurately separate organs, tissues, or target regions from the background in medical images, providing essential auxiliary information for clinical diagnosis, lesion detection, and treatment planning. However, medical images often suffer from significant resolution variations, complex tissue structures, and interference from noise and artifacts, placing high demands on the accuracy of segmentation methods.

[0003] Against the backdrop of the rapid development of deep learning, the U-Net architecture, as a pioneering work in medical image segmentation, has been widely applied and achieved remarkable results. U-Net effectively improves segmentation accuracy through its symmetrical encoder-decoder structure and skip connections. However, the U-Net architecture still has certain limitations: its convolutional structure struggles to capture long-range dependencies, leading to insufficient performance when handling complex scenes and global contextual information. With the successful application of Transformers in natural language processing, researchers began to introduce them into the field of medical image segmentation. TransUnet embeds the Transformer module into the U-Net encoder, leveraging its powerful global modeling capabilities to compensate for the shortcomings of convolutional neural networks in long-range dependency modeling. TransUnet has improved upon the shortcomings of traditional U-Net to some extent and promoted the development of medical image segmentation. Nevertheless, these methods still have two key problems: first, the semantic gap in skip connections, i.e., the significant difference between low-level spatial features and high-level semantic features between the encoder and decoder; second, the loss of local information during the upsampling process, making the decoder insufficient in recovering target boundaries and fine-grained structures.

[0004] This invention proposes a medical image segmentation method based on GM-TransUnet. This method improves the TransUnet model by adding convolutional local attention (CLAttn) and global-local attention (GLAttn) to enhance the feature quality transmitted by skip connections, thereby effectively alleviating the semantic gap problem caused by skip connections. Furthermore, a multi-scale fusion attention module (MFAttn) is added, which significantly reduces the loss of local information during upsampling through a multi-scale feature fusion mechanism. The synergistic effect of these modules effectively improves the network's boundary localization and overall segmentation accuracy. This invention is applicable to clinical medical image segmentation scenarios and has good practical application value and promising prospects for widespread adoption. Summary of the Invention

[0005] The purpose of this invention is to propose a medical image segmentation method based on GM-TransUnet, which aims to solve the problems of semantic gaps in skip connections and loss of local information in upsampling in existing U-shaped networks.

[0006] The technical solution of this invention is to design and train the proposed model, and then test the model on a public medical dataset.

[0007] Medical image segmentation networks based on GM-TransUnet, such as Figure 1 As shown. The framework consists of an encoder, a decoder, and skip connections.

[0008] (1) Encoder

[0009] The encoder structure mainly consists of three parts: convolutional modules, embedding layers, and Transformer modules. The overall architecture includes three convolutional layers. With each convolutional layer, the spatial resolution of the input feature map is halved, while the number of channels doubles. This hierarchical encoding method significantly improves the expressive power of the features. The embedding layer acts as a transition from convolutional features to the Transformer layer, achieving the necessary dimensionality mapping. The Transformer layer uses a multi-head self-attention mechanism to globally model the input features, enabling the network to capture long-range dependencies and contextual semantic information of the image during the feature encoding stage, thereby improving the comprehensiveness of feature representation. This encoder combines the local feature extraction capability of convolution with the global dependency modeling of Transformer, allowing the model to more comprehensively understand image content, thus demonstrating excellent performance in medical image segmentation tasks.

[0010] (2) Jump connection

[0011] Skip connections pass shallow features from the encoder to the decoder, enabling the model to utilize both high-level semantic information and low-level spatial details during the decoder stage. However, significant differences exist between the encoder and decoder at the semantic level. Directly passing encoder features to the decoder via skip connections for fusion can easily lead to semantic gaps and information redundancy, thus reducing the model's segmentation accuracy. To address these issues, we introduce CLAttn and GLAttn modules into the skip connections, effectively suppressing redundant information, highlighting key information, mitigating semantic gaps, and improving the overall quality of the feature maps.

[0012] (3) Decoder

[0013] The decoder structure mainly consists of four core parts: an upsampling convolutional module, an MFAttn module, feature fusion, and a segmentation head. The upsampling convolutional module is primarily used to restore spatial resolution layer by layer. With each upsampling convolutional layer, the feature map size doubles, ultimately restoring deep semantic features to a spatial resolution consistent with the original image. The MFAttn module performs multi-scale attention enhancement on the features after upsampling, emphasizing key region information and reducing local detail loss, thereby effectively improving segmentation accuracy. Feature fusion combines features passed through skip connections with decoder features, achieving a collaborative expression of global semantic information and local details. Finally, the segmentation head maps the final output feature map back to the original image size, generating the final segmentation result. Overall, the decoder structure effectively combines the semantic features output by the encoder with the information passed through skip connections, achieving efficient and precise feature recovery.

[0014] (4) Loss function and evaluation function

[0015] The experiment used a weighted combination of Dice loss and cross-entropy loss as the loss function. Combining these two loss functions can simultaneously ensure pixel-level accuracy and region-level consistency, thereby achieving better performance and a more stable optimization process in real-world scenarios such as class imbalance and blurred boundaries. The calculation formula is shown in Equation 1.

[0016] (1)

[0017] In the model performance evaluation phase, we selected three commonly used evaluation metrics for comprehensive analysis: Intersection over Union (IoU), Dice coefficient, and Precision. These metrics are widely used in medical image segmentation tasks and have become important standards for measuring model performance. Specifically, IoU mainly reflects the degree of overlap between the predicted result and the real region, that is, how much the region segmented by the model overlaps with the labeled region. A higher IoU indicates a better overall match, and the calculation formula is shown in Equation 2. The Dice coefficient is similar to IoU, but it places more emphasis on the consistency between the predicted region and the real region at the boundary and in the overlapping parts. It is more sensitive to small targets or imbalanced data, and therefore is widely used in medical images, and the calculation formula is shown in Equation 3. Precision focuses on how many regions predicted by the model as lesions are actually lesions, reflecting the model's "false positive" situation, and the calculation formula is shown in Equation 4. By using the above metrics comprehensively, the segmentation performance of the model can be evaluated from different perspectives.

[0018] (2)

[0019] (3)

[0020] (4) Attached Figure Description

[0021] Figure 1 Medical image segmentation network based on GM-TransUnet

[0022] Figure 2 CLAttn architecture diagram

[0023] Figure 3 GLAttn Architecture Diagram

[0024] Figure 4 MFAttn Architecture Diagram

[0025] Figure 5 Visualization of CVC-ClinicDB dataset splitting results

[0026] Figure 6 Visualization of Kvasir-Instrument dataset segmentation results

[0027] Figure 7 Model test evaluation index scores Detailed Implementation

[0028] This invention relates to a medical image segmentation method based on GM-TransUnet, and the specific steps are as follows:

[0029] (1) First, we divided the training and testing data. We trained and evaluated on publicly available medical image datasets. The CVC-ClinicDB dataset, provided by the Barcelona Computer Vision Center in Spain, contains 612 high-resolution colonoscopy images, each with a precisely labeled polyp mask. It is mainly used for polyp detection, segmentation, and computer-aided diagnostic systems. For this dataset, our splitting strategy was to randomly select 550 images for training and 62 for testing. The Kvasir-Instrument dataset focuses on the segmentation of diagnostic and therapeutic tools used in gastrointestinal endoscopy, including biopsy forceps, snares, balloon dilators, etc. It contains 590 high-quality endoscopic images, each with its corresponding mask label. For this dataset, our splitting strategy was to randomly select 472 images for training and 118 for testing.

[0030] (2) Build the CLAttn and GLAttn modules. The Convolutional Local Attention (CLAttn) module is as follows: Figure 2As shown, the overall structure consists of three parts: depthwise separable convolution, an efficient local attention mechanism (ELA), and a final residual connection. Depthwise separable convolution significantly reduces computation while still extracting key initial local feature information. The introduced efficient local attention mechanism aggregates global information from the feature map along both the horizontal and vertical directions using strip pooling, then performs one-dimensional convolution, group normalization, and non-linear activation to generate spatial weight vectors in two directions. These two weight vectors are then multiplied and broadcast spatially to the original feature map size, resulting in a two-dimensional spatial attention map. This attention map assigns a fine-grained weight to each location in the feature map, achieving efficient local feature enhancement. Finally, residual connections fuse the original input features with the attention-enhanced features, improving the expressive power of key features while effectively preserving original texture information and avoiding detail loss.

[0031] Since shallow skip connection features only pass through one convolutional block, their semantic expressive power is limited, leading to the most significant semantic gap between them and encoder features. To address this issue, we improve the CLAttn module and propose the Global-Local Attention (GLAttn) module, whose structure is as follows: Figure 3 As shown, this module, based on CLAttn, introduces two additional depthwise separable convolutions and two Transformer blocks. The Transformer's multi-head self-attention (MHSA) mechanism directly establishes dependencies between any two locations, allowing geographically distant regions in the feature map to communicate with each other, thus establishing global consistency at the semantic level. Simultaneously, the self-attention mechanism, through weighted aggregation of features, enables the model to effectively suppress noise interference while strengthening key features.

[0032] For highly semantic deep skip connection features, the original texture information they contain is limited. Therefore, we removed the residual connections in CLAttn to reduce the introduction of redundant information and noise.

[0033] (3) Build the MFAttn module. MFAttn adopts a multi-branch parallel structure, and the detailed architecture is as follows: Figure 4As shown, the first branch captures multi-scale features through parallel combinations of different convolutional kernels, while simultaneously extracting global contextual information using global average pooling. Subsequently, the output of this branch is fused with the output features of the second branch using matrix addition, thus providing global semantic guidance for the overall features. The second branch first divides the input feature map into multiple subgroups for parallel processing to reduce computational cost. Then, each subgroup is further divided into two sub-paths: one sub-path performs one-dimensional global average pooling along both the horizontal and vertical directions to capture channel encoding information in different spatial dimensions. The pooling results are then concatenated, and channel dependencies are established using 1×1 convolutions without dimensionality reduction, ultimately followed by Sigmoid activation; the other sub-path uses dilated convolutions, significantly expanding the receptive field while maintaining a smaller number of parameters, thereby enhancing feature capture capabilities. Finally, the features from the two sub-paths are fused using a cross-spatial information aggregation method, resulting in richer feature representations.

[0034] (4) We built the GM-TransUnet model based on the PyTorch framework and trained it on a single NVIDIA RTX3080ti GPU. The input image resolution was 224×224, the learning rate was set to 1e-2, the momentum was 0.9, and the weight decay was 1e-4. Considering the differences in the size of different datasets, the CVC-ClinicDB dataset was trained for 100 epochs, while the Kvasir-Instrument dataset was trained for 150 epochs. The batch size for both datasets was set to 16.

[0035] (5) The target was segmented using a medical image segmentation method based on GM-TransUnet. Figure 5 Figure 6 presents the visualization results of the model segmentation. Figure 7 It is the score result of the model test evaluation index.

Claims

1. A medical image segmentation method based on GM-TransUnet, characterized in that, Includes the following steps: S1: Prepare training data by dividing the public dataset into training and test sets according to the proportions, and preprocess the input images; S2: The input to this method is the medical image to be segmented; S3: The GM-TransUnet architecture was built in the PyTorch framework, which includes an encoder, skip connections, and a decoder; S4: The encoder architecture consists of three key components: a convolutional module, an embedding layer, and a Transformer module; S5: Skip connections pass shallow features from the encoder to the decoder, enabling the model to utilize both high-level semantic information and low-level spatial detail information in the decoder stage. Introducing convolutional local attention (CLAttn) and global-local attention (GLAttn) modules into skip connections can effectively suppress redundant information, highlight key information, and alleviate the semantic gap, thereby improving the quality of features passed by skip connections. S6: The decoder architecture consists of four key components: upsampling convolutional block, MFAtn, feature fusion part, and segmentation head. After the upsampling operation, a multi-scale fusion attention (MFAttn) module is introduced to enhance the features, highlighting key regional information and reducing the loss of local details, thereby effectively improving the segmentation accuracy. S7: Set the training parameters. The input image resolution is 224×224, the learning rate is 1e-2, the momentum is 0.9, and the weight decay is 1e-4. Considering the size difference of different datasets, the CVC-ClinicDB dataset was trained for 100 epochs, while the Kvasir-Instrument dataset was trained for 150 epochs. The batch size of both datasets was set to 16. The loss function used during model training is calculated by formula (1): (1) S8: Train the model according to the set parameters to obtain the training weights of the entire model; S9: The evaluation metrics used in the model testing are Intersection over Union (IoU), Dice coefficient, and Precision; the specific calculation formulas are shown in equations (2), (3), and (4) below: (2) (3) (4) S10: The GM-TransUnet network is tested using a partitioned test set and the training weights of the entire model. The network performance is evaluated using evaluation metrics. IoU mainly reflects the degree of overlap between the predicted result and the real region, that is, how much the region segmented by the model overlaps with the labeled region. The higher the IoU, the better the overall matching. The Dice coefficient is similar to IoU, but it emphasizes the consistency between the predicted region and the real region at the boundary and the overlapping part. It is more sensitive to small targets or imbalanced data, so it is widely used in medical images. Precision focuses on how many regions predicted by the model as lesions are actually lesions, and can reflect the model's "false positives". By using the above metrics in combination, the segmentation performance of the model can be comprehensively evaluated from different perspectives.

Images