A medical image segmentation method and system based on a double-branch decoder
By combining a CNN-Transformer hybrid encoder and a kMaX Transformer enhanced decoder, the problems of insufficient utilization of category information and insufficient capture of detailed information by the decoder in existing medical image segmentation models are solved, achieving higher accuracy and robustness in medical image segmentation.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-04-03
- Publication Date
- 2026-03-17
AI Technical Summary
Existing medical image segmentation models fail to fully utilize category information when using Transformers, and the decoder fails to maximize the capture of long-range dependencies, resulting in a lack of global consistency and detailed information in the segmentation results.
Image features are extracted using a CNN-Transformer hybrid encoder, combined with a pixel decoder and a kMaX Transformer enhanced decoder. By designing and optimizing model parameters through comparative clustering assignment loss and multi-scale segmentation loss, detailed information can be restored and global features can be captured.
It improves the accuracy and robustness of medical image segmentation, especially when dealing with complex anatomical structures and regions with blurred boundaries, enhancing the model's segmentation precision and generalization ability.
Smart Images

Figure CN118397267B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to a medical image segmentation method and system based on a dual-branch decoder, belonging to the field of image processing technology. Background Technology
[0002] Medical image segmentation plays a crucial role in clinical treatment, providing valuable information for clinical diagnosis and follow-up. Traditional image segmentation methods in clinical medicine rely heavily on rigorous manual intervention algorithms and specialized knowledge. However, this process is time-consuming, costly, and can lead to physician fatigue, thus affecting diagnostic accuracy. With advancements in computer hardware, deep learning-based image segmentation technology has gained widespread attention. This technology can accurately segment anatomical structures and pathological regions from medical images, improving patient treatment outcomes. Medical image segmentation technology has become a key tool for improving the efficiency and accuracy of medical consultations.
[0003] Among numerous deep learning-based segmentation methods, neural networks based on CNN architectures, especially fully convolutional networks (FCNs), dominate medical image segmentation due to their ability to handle images of arbitrary sizes. U-Net and its variants, in particular, have attracted significant attention for their effectiveness in capturing spatial context and generating accurate segmentation masks. These CNN-based methods extract feature information by stacking convolutional and pooling layers. Convolutional layers extract features within local regions using convolutional kernels, and then pooling layers reduce the spatial size of the feature maps, decreasing computational complexity and the number of parameters. However, this approach may neglect global contextual information, leading to a lack of global consistency in segmentation results, especially challenging when dealing with complex anatomical structures.
[0004] To address the limitations of traditional CNN-based models in handling long-range dependencies, the Transformer was proposed. By introducing a self-attention mechanism, it models global dependencies, achieving significant progress in Natural Language Processing (NLP) due to its independence from sequence length limitations. Subsequently, researchers began exploring the application of the Transformer in the field of vision. The Vision Transformer (ViT) was the first to apply the Transformer to image recognition tasks, achieving remarkable success and demonstrating its effectiveness and potential in image processing. Despite this success, the Transformer faces challenges in effectively handling spatial information and adapting to the pixel-level nature of medical images.
[0005] Integrating the Transformer architecture into medical image segmentation models has sparked the development of novel hybrid approaches. A pioneer in hybrid models for medical image processing is the TransUNet network, which fully leverages the strengths of Convolutional Neural Networks (CNNs) and Transformers. Specifically, it utilizes ResNet50 to extract low-level features, then encodes these features through a Transformer to capture long-range correlations. By stacking CNNs and Transformers, TransUNet effectively learns detailed image features and global background simultaneously, thereby improving segmentation accuracy and generalization performance. However, previous work primarily focused on integrating the strengths of Transformers and CNNs at the encoder stage, without making any improvements to the decoder, neglecting the importance of optimizing the decoder. In the field of medical image segmentation, we often face the challenge of imbalanced sample datasets. Existing pixel-level segmentation models typically classify each pixel through feature extraction and processing within the network to generate segmentation results. However, these models do not fully utilize category information to assist the segmentation task, resulting in a certain impact on the model's prediction performance. Mask Transformers attempts to enhance the CNN-based backbone network through independent Transformer blocks. MaXDeeplab interprets object queries in DETR as memory-encoded queries for end-to-end panoptic segmentation. MaxQuery and kMaX-Deeplab propose interpreting queries as cluster centers and adding regulatory constraints to learn the clustering representation of the queries. To our knowledge, there is currently limited research on Mask Transformer-based methods, and there is still significant room for research on how to better integrate Mask Transformer into existing backbone networks. Summary of the Invention
[0006] This invention addresses the technical problems of existing U-shaped network decoders, which typically use simple cascaded upsampling and fail to maximize the Transformer's ability to capture long-range dependencies, and existing pixel-level segmentation models that do not utilize category information. Therefore, this invention proposes a medical image segmentation method and system based on a dual-branch decoder.
[0007] The technical solution adopted by this invention to solve the above problems is as follows: This invention proposes a medical image segmentation method based on a dual-branch decoder, comprising:
[0008] Step 1: Extract image features from medical images based on a CNN-Transformer hybrid encoder;
[0009] Step 2: Construct a pixel decoder and restore the resolution of the feature map of the medical image based on the pixel decoder;
[0010] Step 3: Construct an enhanced decoder containing an LMaX Transformer, and perform clustering based on the enhanced decoder to capture detailed information of the decoded medical image features and obtain the segmentation prediction results of the medical image;
[0011] Step 4: Design a comparative clustering assignment loss, obtain a feature augmentation view, and enhance the similarity between the feature augmentation view and the clustering results;
[0012] Step 5: Design a multi-scale segmentation loss. Obtain the total segmentation loss based on the contrastive clustering assignment loss and the multi-scale segmentation loss. Optimize the parameters of the hybrid encoder and the enhanced decoder through backpropagation of the total segmentation loss to complete the automatic segmentation of medical images.
[0013] Optionally, step 1, which involves extracting image features from medical images, includes:
[0014] Step 1.1: Convert the medical image x∈R H×W×C The input is fed into the CNN encoder for scale division, and then... and Feature maps at four scales, with each scale's output hidden feature sequence being... Where C is the number of channels, H is the image height, and W is the image width;
[0015] Step 1.2: Output the hidden feature sequence from the feature maps at each scale in the CNN encoder. Reshape to obtain Image blocks Each image block is P×P in size;
[0016] Step 1.3: Use linear projection to map the reshaped image patches into D-dimensional embeddings and record the position information of each image patch to obtain the local features of the medical image;
[0017] Step 1.4: Input the D-dimensional embedding into a Transformer encoder composed of multi-head self-attention and multilayer perceptrons to obtain global features of the medical image. The hidden feature sequence output by the CNN-Transformer hybrid encoder is as follows:
[0018] Optionally, step 2, which involves restoring the resolution of the feature map of the medical image, includes:
[0019] Step 2.1: Output the hidden feature sequence from the hybrid encoder. Reshaping at the scale yields a scale of... The hidden feature sequence;
[0020] Step 2.2: By cascading multiple upsampling blocks, the resolution of the scale-reshaped hidden feature sequence is increased from... Gradually recover to H×W.
[0021] Optionally, step 3, which involves obtaining the segmentation prediction results of the medical image, includes:
[0022] Step 3.1: Fuse pixel features from the pixel decoder and the CNN-Transformer hybrid encoder at four scales via over-hop connections;
[0023] Step 3.2: Input the pixel features of each scale after fusion into the kMaX Transformer module for clustering;
[0024] Step 3.3: Initialize the cluster centers using a truncated normal distribution to obtain the output cluster centers B0, B0∈R. N ×C ;
[0025] Step 3.4: Update the cluster centers;
[0026] Step 3.5: Aggregate the pixels within each cluster based on the updated cluster centers to obtain the clustering results;
[0027] Step 3.6: Obtain the segmentation prediction results of the medical image based on the clustering results, and output the feature map f.
[0028] Optionally, the step of updating the cluster centers in step 3.4 includes:
[0029] Step 3.4.1: Set the scale obtained from the first upsampling block to... The pixel features and cluster center B0 are used as input to the first-stage kMaX decoder to obtain the first-stage updated cluster center B1.
[0030] Step 3.4.2: Sequentially use the pixel features output from all upsampled blocks and the cluster centers obtained in each stage as input to the kMaX decoder in the next stage to complete the update of the cluster centers.
[0031] Optionally, the steps in step 3.4.2 for obtaining the cluster centers obtained at each stage include:
[0032] Step 3.4.2.1: Convolve and normalize the pixel features and cluster centers respectively, and fuse the pixel and cluster center information using k-means cross attention;
[0033] Step 3.4.2.2: Use the multi-head attention module to aggregate and optimize the fused pixel and cluster center information, and obtain the cluster centers;
[0034] Step 3.4.2.3: Update the cluster centers output by the multi-head attention module based on the feedforward network to obtain the updated cluster centers at each stage;
[0035] The expression for the k-means cross-attention operation is:
[0036] z m =BN(Conv(z) m ))(1)
[0037] B l "=KCA(BN(Conv(B l-1 )),z m )+B l-1 (2)
[0038] In formulas (1) and (2), Conv() is a convolutional layer, BN(·) is a batch normalization operator, KCA() is k-means cross-attention, and B l-1 B serves as the cluster center for the next higher level. l "This is the output of the l-th KCA module;
[0039] The expression for the multi-head attention mechanism operation is:
[0040] B l =MSA(BN(Conv(B l ”)))+B l (3)
[0041] In formula (3), MSA is a multi-head self-attention module;
[0042] The expression for updating cluster centers is:
[0043] B l =FFN(BN(Conv(B l ')))+B l (4)
[0044] In formula (4), FFN is the feedforward neural network module, B l This is the output of the first-level FFN module.
[0045] Optionally, the step of aggregating pixels within each cluster in step 3.5 includes:
[0046] Step 3.5.1: Stack multiple layers of kMaXTransformer modules to form an enhanced decoder;
[0047] Step 3.5.2: Aggregate the pixels within each cluster based on the cross-attention mechanism of the kMaXTransformer module and the updated cluster centers to obtain the cluster assignment results;
[0048] The expression for calculating cross attention is:
[0049]
[0050] In formula (5), C and P are the query and pixel features, respectively, and Q is the pixel feature. c ,K p V p These are linear projection queries, key-value pairs, and the aggregation expression is:
[0051]
[0052]
[0053] In formulas (6) and (7), B, P, and A represent the cluster center, pixel feature, and cluster assignment, respectively, and B∈R N×C ,P∈R HW×C ,A∈R HW×N .
[0054] Optionally, step 4, which involves obtaining a feature augmentation view and enhancing the similarity between the feature augmentation view and the clustering results, includes:
[0055] Step 4.1: Add noise to the feature map f output by the decoder to generate two enhanced views f1 and f2;
[0056] Step 4.2: Maximize the similarity between the feature map and the cluster classification based on the enhanced views f1 and f2 to obtain the comparative clustering loss;
[0057] The expression for the clustering loss is as follows:
[0058]
[0059] In formula (8), τ is the temperature parameter, f1 and f2 are both feature-enhanced views, B is the clustering classification result, and q1 and q2 are calculated from the clustering classification result B and the feature-enhanced views f1 and f2.
[0060] Optionally, step 5, which involves inversely optimizing the parameters of the hybrid encoder and the enhanced decoder using the total loss, includes:
[0061] Step 5.1: Downsample the Ground Truth in the hybrid encoder and the enhancement encoder to obtain... Ground Trurh at four scales: H×W;
[0062] Step 5.2: Calculate the multi-scale segmentation loss of the feature maps at different scales and the corresponding Ground Truth at the output of the enhanced decoder containing the kMaX Transformer.
[0063] Step 5.3: Weight the multi-scale segmentation loss and the contrastive clustering loss to obtain the total segmentation loss;
[0064] Step 5.4: Optimize the parameters of the hybrid encoder and the enhanced decoder by back-optimizing the total segmentation loss to complete the automatic segmentation of the medical image;
[0065] The expression for multi-scale segmentation is:
[0066]
[0067] In formula (9), L ce and L dice These are the binary cross-entropy loss and the Dice loss, respectively. λ is a hyperparameter used to balance the segmentation loss per pixel, and λ1, λ2, λ3, and λ4 are constants, with λ1, λ2, λ3, and λ4 being 0.05, 0.05, 0.1, and 0.8, respectively.
[0068] The expression for the total segmentation loss is:
[0069] L total =αL seg +βL ccl (10)
[0070] In formula (10), the parameters α and β are constants, with α being 0.5 and β being 0.1 respectively.
[0071] A medical image segmentation system based on a hybrid encoder and a dual-branch decoder includes:
[0072] The medical image feature acquisition module is used to extract image features from medical images;
[0073] The medical image resolution restoration module is used to restore the resolution of feature maps in medical images.
[0074] The medical image segmentation and prediction module is used to obtain the segmentation and prediction results of medical images.
[0075] The total segmentation loss acquisition module is used to acquire the exchange-contrast clustering assignment loss and the multi-scale segmentation loss, and to acquire the total segmentation loss based on the exchange-contrast clustering assignment loss and the multi-scale segmentation loss.
[0076] The automatic medical image segmentation module is used to inversely optimize the parameters of the hybrid encoder and the enhancement decoder by using the total segmentation loss to complete the automatic segmentation of medical images.
[0077] The beneficial effects of this invention are:
[0078] 1. This invention proposes a dual-branch decoder, a pixel decoder, and a kMaX Transformer augmented decoder, building upon the traditional U-shaped network for image segmentation. Compared to the pixel decoder with only simple upsampling, the augmented decoder fully utilizes the Transformer's ability to capture long-range dependencies, avoiding the loss of detailed information. We learn the cross-attention between cluster centers and pixel features, using the cluster assignment concept to achieve segmentation, compensating for the detailed information lost during upsampling and obtaining more accurate segmentation results. Experimental results on heart and multi-organ datasets fully demonstrate the superior performance of kMaX-TransU.
[0079] 2. This invention proposes an exchange-contrast clustering assignment loss, which improves the consistency between feature-enhanced views by applying clustering assignment to enhanced features, thereby enhancing the decoding capability of the decoder and further improving the accuracy and robustness of segmentation.
[0080] 3. This invention proposes a multi-scale segmentation loss, which uses the four-dimensional features output by the model decoder and their true values as the segmentation loss, enabling the model to better learn the feature representations of different regions in the image and to be more robust when dealing with structures of different scales or semantic levels, thereby improving the model's generalization ability. Attached Figure Description
[0081] Figure 1 This is a flowchart of a medical image segmentation method based on a hybrid encoder and a dual-branch decoder provided in an embodiment of the present invention;
[0082] Figure 2 This is a structural diagram of the hybrid encoder, pixel decoder, and enhancement decoder provided in an embodiment of the present invention;
[0083] Figure 3 This is a structural diagram of the kMaX Transformer Block provided in an embodiment of the present invention;
[0084] Figure 4 A flowchart illustrating the enhanced decoder, contrastive clustering assignment loss, and multi-scale segmentation loss provided in this embodiment of the invention;
[0085] Figure 5 Segmentation results of various methods for the Synapse multi-organ CT dataset provided in this embodiment of the invention;
[0086] Figure 6A comparison chart of segmentation results of various methods on the ACDC dataset provided in this embodiment of the invention;
[0087] Figure 7 A flowchart of a medical image segmentation system based on a hybrid encoder and a dual-branch decoder is provided for an embodiment of the present invention;
[0088] In the diagram, CNN stands for CNN neural network, Transformer is the attention mechanism, kMaX Transformer is the k-means Mask Transformer network, C is the number of channels, H is the image height, W is the image width, N is the number of classes, K is the number of clusters, Ground Truth is the ground truth, i.e., the segmentation map annotated by the doctor, Hybrid encoder is the hybrid encoder, Pixeldecoder is the pixel decoder, Enhanced decoder is the enhanced decoder, CNN stands for convolution, Transformer is the transformer, hidden features are the latent features, cluster centers are the cluster centers, contrastive cluster loss is the contrastive cluster loss, Multi-scale segmentation loss is the multi-scale segmentation loss, downsample is the downsampling method, upsample is the upsampling method, segmentation head is the segmentation head, Feature Fusion is the feature fusion method, ConvBN is the convolutional batch normalization method, k-means cross attention is the k-means cross attention method, multi-head self-attention is the multi-head self-attention method, feed-forward network is the feed-forward network, and cluster-wise is the cluster-wise network. argmax is the argmax of the cluster level, linear is the linear layer, and aug is the enhancement layer. Example
[0089] Combination Figure 1-4 and Figure 7 This embodiment will be described as follows: Figure 1 and Figure 2 As shown, the steps of the medical image segmentation method based on a hybrid encoder and a dual-branch decoder proposed in this invention include:
[0090] Step S1: Design a CNN-Transformer hybrid dual-branch decoder to extract local and global features of the image;
[0091] The structure of the CNN-Transformer hybrid encoder is as follows: Figure 2 As shown on the left;
[0092] Furthermore, step S1 includes:
[0093] Step S101: Convert the medical image x∈R H×W×C The inputs are fed into the CNN encoder, and the results are obtained respectively. Feature maps at four scales;
[0094] Step S102: Extract the latent features from the final output of the CNN encoder. Remodeled into a series of flat 2D patches Each patch is P×P in size, and the number of image patches is...
[0095] Step S103: Map the vectorized block Xp to a D-dimensional embedding using linear projection. Furthermore, we add a position embedding to the patch embedding to record position information:
[0096] Step S104: Input into the Transformer encoder, which consists of an L-layer multi-head self-attention (MSA) and a multilayer perceptron (MLP).
[0097] Step S2: Based on the upsampled pixel decoder, perform preliminary decoding of the features from step S1;
[0098] Furthermore, step S2 includes:
[0099] Step S201: Analyze the hidden feature sequence output by the hybrid encoder. Reshape into shape Format;
[0100] Step S202: Increase the resolution from [previous step] by cascading multiple upsampling blocks. Gradually recover to H×W;
[0101] Step S203: Fuse features at four scales of the pixel decoder and hybrid encoder by skipping connections.
[0102] Step S3: Design an enhanced decoder containing kMaXTransformer to capture detailed information of the features decoded in step S2;
[0103] The enhanced decoder is composed of multiple stacked kMaXTransformer modules;
[0104] The enhanced decoder includes Stage 0, Stage 1, Stage 2, and Stage 3.
[0105] The feature maps at four different scales are decoded in four stages.
[0106] Furthermore, step S3 includes:
[0107] Step S301: Initialize cluster centers using a truncated normal distribution to obtain initial cluster centers C0∈R. N×C ;
[0108] Step S302: Convolve and batch normalize the cluster centers and corresponding upsampled block pixel features of each stage, respectively. Then, fuse the pixel and cluster center information using k-means cross-attention. This can be represented as:
[0109] z m =BN(Conv(z) m ))(1)
[0110] B l "=KCA(BN(Conv(B l-1 )),z m )+B l-1 (2)
[0111] In formulas (1) and (2), Conv() is a convolutional layer, BN(·) is a batch normalization operator, KCA() is k-means cross-attention, and B l-1 B serves as the cluster center for the next higher level. l "This is the output of the l-th KCA module;
[0112] Step S303: Further aggregate and optimize information using multi-head self-attention, which can be represented as:
[0113] B l =MSA(BN(Conv(B l ”)))+B l (3)
[0114] In formula (3), MSA is a multi-head self-attention module;
[0115] Step S304: Update the cluster centers output by the MSA module using a feedforward network (FFN). The update expression is:
[0116] B l =FFN(BN(Conv(B l ')))+B l (4)
[0117] In formula (4), FFN is the feedforward neural network module, B l This is the output of the l-th layer FFN module;
[0118] Step S305: Aggregate the pixels within each cluster based on the updated cluster centers to obtain cluster assignments, which can be represented as:
[0119]
[0120]
[0121] In formulas (5) and (6), B∈R N×C ,P∈R HW×C ,A∈R HW×N These represent cluster centers, pixel features, and cluster assignments, respectively.
[0122] Furthermore, step S302 includes:
[0123] Step S30201: The cross-attention mechanism aggregates pixels within each cluster. The calculation expression for cross-attention is:
[0124]
[0125] In formula (7), the superscripts c and p represent the query and pixel features, respectively. c K p V p Represents linear projection queries, keys, and values.
[0126] Step S4: Design a comparison between clustering assignment loss and multi-scale segmentation loss, optimize the segmentation results through backpropagation, and finally achieve automatic image segmentation.
[0127] The total loss is obtained by adding the contrastive clustering assignment loss and the multi-scale segmentation loss, which together optimize the model parameters.
[0128] Furthermore, step S4 includes:
[0129] Step S401: Design a comparative clustering assignment loss. By applying clustering assignment to augmented features, the consistency between feature-augmented views can be improved.
[0130] Step S402: Design a multi-scale segmentation loss, while focusing on and utilizing image features and label information at different scales, thereby improving the accuracy and robustness of medical image segmentation;
[0131] Step S403: By controlling the weight parameters, the clustering assignment loss and multi-scale segmentation loss are compared and work together to optimize the model parameters. The expression for the total loss is:
[0132] L total =αL seg +βL ccl (8)
[0133] In formula (8), the parameters α and β are constants, with α being 0.5 and β being 0.1, respectively;
[0134] Furthermore, step S401 includes:
[0135] Step S40101: Add noise to the feature map f output by the decoder to generate two enhanced views f1 and f2;
[0136] Step S40102: Maximize the similarity between the feature map and the cluster classification based on the enhanced views f1 and f2;
[0137] The expression for the clustering assignment loss is as follows:
[0138]
[0139] In formula (9), τ is the temperature parameter, f1 and f2 are both feature-enhanced views, B is the clustering classification result, and q1 and q2 are calculated from the clustering classification result B and the feature-enhanced views f1 and f2.
[0140] Furthermore, step S402 includes:
[0141] Step S40201: Downsample the Ground Truth to obtain Ground Trurh at four scales: H×W;
[0142] Step S40202: Perform segmentation loss on the feature maps of different granularities output by the kMaX decoder and the ground truth of the corresponding scale;
[0143] The expression for the multi-scale segmentation loss is:
[0144]
[0145] In formula (10), L ce and L dice These are the binary cross-entropy loss and the Dice loss, respectively. λ is a hyperparameter used to balance the segmentation loss per pixel, and λ1, λ2, λ3, and λ4 are constants, with λ1, λ2, λ3, and λ4 being 0.05, 0.05, 0.1, and 0.8, respectively.
[0146] like Figure 7 As shown, the present invention proposes a medical image segmentation system based on a hybrid encoder and a dual-branch decoder, comprising:
[0147] The medical image feature acquisition module is used to extract image features from medical images;
[0148] The medical image resolution restoration module is used to restore the resolution of feature maps in medical images.
[0149] The medical image segmentation and prediction module is used to obtain the segmentation and prediction results of medical images.
[0150] The total segmentation loss acquisition module is used to acquire the exchange-contrast clustering assignment loss and the multi-scale segmentation loss, and to acquire the total segmentation loss based on the exchange-contrast clustering assignment loss and the multi-scale segmentation loss.
[0151] The automatic medical image segmentation module is used to inversely optimize the parameters of the hybrid encoder and the enhancement decoder by using the total segmentation loss to complete the automatic segmentation of medical images.
[0152] By employing multi-scale segmentation loss, we are able to more comprehensively consider the structure and details at different scales in the image, thereby better adapting to targets and lesions at different scales.
[0153] Example 2
[0154] Combination Figure 5-6 This embodiment will be described, and the model of this invention will be compared with the following five models:
[0155] 1) UNet: The structure of the UNet model includes a downsampling path and an upsampling path, which are combined in a shape similar to the letter U. Downsampling captures the contextual information of the image and gradually reduces the spatial dimension through successive convolution and max pooling layers; upsampling gradually restores the spatial dimension of the image, while retaining positional information by combining the corresponding feature maps in the shrinking path through skip links.
[0156] 2) TransUnet: TransUnet is the first architecture to apply Transformer to medical image segmentation. It uses Transformer as an encoder to extract global features, and CNN as a decoder to perform accurate local feature reconstruction and upsampling.
[0157] 3) Swin-Unet: Swin-Unet adapts to different image scales by using a hierarchical Transformer and introduces a shift window mechanism to improve computational efficiency, capture long-distance dependencies and rich contextual information.
[0158] 4) LeViT-Unet: LeViT is used as an encoder to capture rich global features, while the UNet structure is responsible for fine local feature recovery and upsampling.
[0159] 5) TransCeption: TransCeption is a U-shaped structure that improves the Transformer through multiple scales. The encoder can capture multi-scale representations within a single stage, constructing an intra-stage feature fusion module to enhance the aggregated feature map from all branches.
[0160] The above models are all based on U-shaped structures and adopt the ideas of CNN and Transformer. This invention utilizes the clustering idea of k-means Mask Transformer to better model the spatial relationship between pixels, improve the segmentation model's ability to understand image structure, and thus improve the accuracy and performance of the segmentation model, especially when processing medical images with complex structures or blurred boundaries.
[0161] Multi-organ dataset: The public multi-organ segmentation dataset Synapse collected 30 abdominal CT scan cases in the 2015 MICCAI challenge. Each CT scan has 18 cases as the training set and 12 cases as the test set. Each case contains 85 to 198 slices. This invention evaluates 8 organs including spleen, liver, aorta, gallbladder, left kidney, right kidney, stomach and pancreas on Synapse.
[0162] The ACDC dataset for cardiac diagnosis consists of 100 patient cases, including 70 training samples, 10 validation samples, and 20 test samples. Each sample contains the left ventricle (LV), right ventricle (RV), and myocardium (MYO) with corresponding annotations.
[0163] The above two datasets were implemented using kMaX-TransU and other models, with DSC and HD95 as evaluation metrics.
[0164] DSC (Dice Similarity Coefficient) The Dice coefficient is a set similarity measurement function, usually used to calculate the similarity between two samples, with a value range between [0,1].
[0165]
[0166] HD (Hausdorff Distance) is a distance metric used to measure the similarity or difference between two sets of points, and it is widely used in fields such as image segmentation. HD95 is an extension of Hausdorff Distance, which refers to the Hausdorff Distance that minimizes the distance between 95% of point pairs.
[0167] HD95 = max k95%[d(X,Y),D(Y,X)]
[0168] The comparison results of each model on different datasets are as follows:
[0169] 1) Comparison results on Synapse
[0170] Table 1 shows the comparison results of kMaX-TransU and other methods on the Synapse multi-organ CT dataset. To provide a more comprehensive evaluation, this invention uses two metrics, DSC and HD95, to evaluate the model's performance. Experimental results show that the proposed method achieves the best performance, with segmentation accuracy as high as 83.15% (DSC↑) and 21.54mm (HD95↓). kMaX-TransU outperforms existing models in DSC and demonstrates excellent overall segmentation performance, proving its strong robustness and reliability in multi-organ segmentation tasks. Furthermore, it can be seen that our proposed method exhibits stronger capabilities in segmenting small organs such as the pancreas. Specifically, the DSC value of the pancreas is 29.97% higher than the V-Net-based result, indicating that our hybrid encoder can effectively extract local and global information from the image and accurately capture detailed information through the enhanced decoder of the dual-branch decoder. Therefore, our kMaX-TransU is more accurate and sensitive in handling fine structures and edge features.
[0171] 2) Comparison results on ACDC
[0172] On the ACDC MRI dataset in Table 2, kMaX-TransU demonstrated the most accurate segmentation results. Specifically, it achieved an overall accuracy of 92.13% on the DSC metric, and also achieved best performance in each individual category (LV, RV, and MYO). The kMaX-TransU model, incorporating an enhanced decoder, effectively utilizes clustering information and pixel features to obtain superior segmentation results. Especially for the less sharp right ventricle, this model generates more ideal segmentation results compared to other methods.
[0173] Experiments further demonstrate that kMaX-TransU with an enhanced decoder improves the model's performance by learning the cross-attention between cluster centers and pixel features and using the concept of cluster assignment to achieve segmentation, thus compensating for the details lost during upsampling. Furthermore, the excellent performance of kMaX-TransU on two datasets proves its generalization and versatility.
[0174] Table 1: Segmentation results on the Synapse dataset
[0175] Method DSC↑ HD95↓ Aorta Gallbladder Kidney (L) Kidney(R) Liver Pancreas Spleen Stomach V-net 68.81 — 75.34 51.87 77.1 80.75 87.84 40.05 80.56 56.98 DARR 69.77 — 74.74 53.77 72.31 73.24 94.08 54.18 89.9 45.96 Unet 76.85 39.7 89.07 69.72 77.77 68.6 93.43 53.98 86.67 75.58 AttnUnet 77.7 36.02 89.55 68.88 77.98 71.11 93.57 58.04 87.30 75.75 TransUnet 77.48 31.69 87.23 63.13 81.87 77.02 94.08 55.86 85.08 75.62 Swin-Unet 79.13 24.62 85.47 66.53 83.28 79.61 94.29 56.58 90.66 76.6 MT-Unet 78.08 26.59 85.87 61.38 84.83 79.36 94.28 57.65 87.74 73.55 LeVit-Unet 78.53 27.32 87.33 62.23 84.61 80.25 93.11 59.07 88.86 72.76 HiFormer 80.39 14.70 86.21 65.69 85.23 79.77 94.61 59.52 90.99 81.08 MORSE 80.85 26.61 88.92 67.53 84.83 81.68 96.83 59.70 87.73 79.58 MISSFormer 81.96 18.20 86.99 68.65 85.21 82.00 94.41 65.67 91.92 80.81 TransCeption 82.24 20.89 87.60 71.82 86.23 80.29 95.01 65.27 91.68 80.02 kMaX-TransU 83.15 21.54 89.81 70.97 86.75 82.56 94.73 70.02 90.33 80.04
[0176] Table 2: Segmentation results of each model on ACDC
[0177] Method DSC (%) ↑ RV Myo LV Unet 87.60 84.62 84.52 93.68 AttnUnet 86.90 83.27 84.33 93.53 TransUnet 89.71 86.67 87.27 95.18 Swin-Unet 88.07 85.77 84.42 94.03 MT-Unet 90.43 86.64 89.04 95.62 LeVit-Unet 90.32 89.55 87.64 93.76 UNETR 88.61 85.29 86.52 94.02 HiFormer 90.40 88.24 87.63 95.30 MISSFormer 91.19 89.85 88.38 95.34 nnUNet 91.61 90.24 89.24 95.36 TransCeption 88.47 87.88 82.87 94.66 nnFormer 91.78 90.22 89.53 95.59 KMax-TransU 92.13 90.76 89.70 95.93
[0178] To ensure the reliability of the experimental results, this invention conducted a detailed analysis and comparison of the two datasets and explored the segmentation effects of four different methods.
[0179] like Figure 5 As shown in the first row, it is evident that other methods have significant shortcomings in identifying the gallbladder, while the method of this invention is significantly superior in segmenting this region. From the third row, it can be seen that other methods fail to accurately infer the pancreatic region when predicting the pancreas, while the method of this invention demonstrates a more accurate predictive ability. Figure 6 As shown, these methods can accurately segment the easily segmented LV region. However, for the more difficult Myo and RV regions, other methods cannot segment them accurately. More specifically, the TransCeption and SwinUNet methods in the second row are relatively coarse in segmenting the Myo boundary, while the TransCeption, TransUNet, and SwinUNet methods in the third row show obvious undersegmentation in RV. kMaX-TransU, on the other hand, demonstrates accurate segmentation capabilities in all three regions. In conclusion, kMaX-TransU exhibits superior segmentation performance on both datasets, fully demonstrating the effectiveness of the hybrid encoder and dual-branch decoder in the model.
[0180] This invention conducted effective ablation experiments on each proposed module to verify its necessity and effectiveness. Using the TransUnet model as a benchmark, a series of experiments were performed on the ACDC dataset, with DSC used as the evaluation metric. Experimental results show that adding the innovative features of this invention to TransUnet significantly improves performance. Specific data can be found in Table 3. First, this invention concatenates an enhanced decoder (kMaX decoder) after the traditional pixel decoder. The results show that the segmentation accuracy reached 91.50%, a 1.79% improvement over the traditional TransUNet segmentation accuracy. Second, we designed a contrastive clustering assignment loss, combining the ideas of contrastive learning and clustering, further improving the segmentation accuracy by 0.45% on top of the enhanced decoder. Furthermore, based on the enhanced decoder, we proposed a multi-scale segmentation loss, enabling the model to more comprehensively utilize the feature information in the image. Although the result showed a slight improvement, the increase was very limited, only 0.15%. However, when we combined these two losses, the segmentation accuracy reached 92.13%, and the DSC was improved by 0.63% compared to the original method with only single-scale segmentation loss. These experimental results strongly demonstrate that in the bi-branch decoder, the enhancement decoder establishes a direct link between classification features and segmentation prediction, resulting in more accurate segmentation predictions. Furthermore, by utilizing the multi-scale information and clustering concepts of the bi-branch decoder, the contrastive clustering assignment loss enhances the model's ability to identify small organs or lesion regions. The multi-scale segmentation loss allows the model to simultaneously focus on and utilize image features and label information at different scales, thereby improving the accuracy and robustness of medical image segmentation.
[0181] Table 3: Roles of each module in the kMaX-TransU method
[0182] Method DSC↑ HD95↓ RV Myo LV Base(TransUnet) 89.71 1.08 86.67 87.27 95.18 Base+kMaXDecoder 91.50 2.27 89.58 89.16 95.77 <![CDATA[Base+kMaXDecoder+L ccl ]]> 91.95 1.10 91.17 88.99 95.70 <![CDATA[Base+kMaXDecoder+L seg ]]> 91.65 1.73 90.08 89.22 95.64 <![CDATA[Base+kMaXDecoder+L ccl +L seg (bears)]]> 92.13 1.07 90.76 89.70 95.93
[0183] The above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention in any way. Although the present invention has been disclosed above with reference to preferred embodiments, it is not intended to limit the present invention. Any person skilled in the art can make some modifications or alterations to the above-disclosed technical content to create equivalent embodiments without departing from the scope of the present invention. Any simple modifications, equivalent substitutions, and improvements made to the above embodiments without departing from the scope of the present invention, based on the technical essence of the present invention and within the spirit and principles of the present invention, shall still fall within the protection scope of the present invention.
Claims
1. A medical image segmentation method based on a two-branch decoder, characterized by, The steps of the medical image segmentation method based on the dual-branch decoder include: Step 1: extracting image features of a medical image based on a CNN-Transformer hybrid encoder; The step of extracting image features of the medical image in step 1 includes: Step 1.1: input the medical image into the CNN encoder for scale division, respectively obtaining feature maps of four scales, and the hidden feature sequence output by each scale feature map is , , and , wherein C is the number of channels, H is the image height, W is the image width; Step 1.2: Hidden feature sequence output for each scale of feature map in CNN encoder Reshaping to get image patches , where each image patch has size ; Step 1.3: mapping the reshaped image block into a D-dimensional embedding using linear projection, and recording the position information of each image block to obtain local features of the medical image; Step 1.4: The D-dimensional embedding is input into a Transformer encoder consisting of multi-head self-attention and multi-layer perceptron to obtain global features of the medical image. The hidden feature sequence output by the CNN-Transformer hybrid encoder is ; Step 2: constructing a pixel decoder, and restoring the resolution of the feature map of the medical image based on the pixel decoder; The step of restoring the resolution of the feature map of the medical image in step 2 includes: Step 2.1: Reshaping the scales of the mixed encoder output hidden feature sequence to get a hidden feature sequence with scales Step 2.2: Reshaping the scales of the mixed encoder output hidden feature sequence to get a hidden feature sequence with scales Step 2.2: The resolution of the scale-remodeled hidden feature sequence is restored from to , , , ; Step 3: constructing an enhanced decoder containing an LMaX Transformer, and capturing detailed information of the decoded medical image features based on the enhanced decoder to obtain a segmentation prediction result of the medical image; The step of obtaining the segmentation prediction result of the medical image in step 3 includes: Step 3.1: fusing pixel features of four scales of the pixel decoder and the CNN-Transformer hybrid encoder through a skip connection; Step 3.2: sequentially inputting the fused pixel features of each scale into a kMaX Transformer module for clustering; Step 3.3: Initialization of cluster centers using truncated normal distribution to get output cluster centers , ; Step 3.4: updating the cluster centers; The step of updating the cluster centers in step 3.4 includes: Step 3.4.1 : The first up-sampling block takes the pixel features of scale as input and produces the first stage updated cluster centers Step 3.4.2: sequentially taking the pixel features output by all up-sampling blocks and the cluster centers obtained at each stage as inputs of a next-stage kMaX decoder to complete the update of the cluster centers; Step 3.5: aggregating pixels in each cluster based on the updated cluster centers to obtain a clustering result; Step 3.6: obtaining a segmentation prediction result of the medical image based on the clustering result, and outputting a feature map f; Step 4: designing a contrastive clustering assignment loss, obtaining a feature enhancement view, and enhancing the similarity between the feature enhancement view and the clustering result; The step of obtaining the feature enhancement view and enhancing the similarity between the feature enhancement view and the clustering result in step 4 includes: Step 4.1: Add noise to the final output feature map f of the decoder to generate two enhanced views and ; Step 4.2: Based on the enhanced view and Maximizing the similarity between the feature map and the cluster classification obtains the contrastive cluster loss; The expression of the contrastive clustering loss is: (8) In formula (8), is a temperature parameter, , are feature enhanced views, is a clustering classification result, and are obtained by clustering the classification result and the feature enhanced views and ; Step 5: designing a multi-scale segmentation loss, obtaining a total segmentation loss based on the contrastive clustering assignment loss and the multi-scale segmentation loss, and optimizing the parameters of the hybrid encoder and the enhanced decoder through the total segmentation loss to complete automatic segmentation of the medical image; The step of optimizing the parameters of the hybrid encoder and the enhanced decoder through the total loss in step 5 includes: Step 5.1: Downsample the Ground Truth in the hybrid encoder and the enhancement encoder to get , , , Ground Truth at four scales; Step 5.2: calculating a multi-scale segmentation loss of feature maps of different scales output by the enhanced decoder containing the kMaX Transformer and the corresponding scale Ground Trurh; Step 5.3: weighting and summing the multi-scale segmentation loss and the contrastive clustering loss to obtain a total segmentation loss; Step 5.4: optimizing the parameters of the hybrid encoder and the enhanced decoder through the total segmentation loss to complete automatic segmentation of the medical image; The expression of the multi-scale segmentation is: (9) In formula (9), and are binary cross-entropy loss and dice loss, respectively, is a hyperparameter for balancing the per-pixel segmentation loss, is a constant, are 0.05, 0.05, 0.1, 0.8, respectively. The expression of the total segmentation loss is: (10) In Equation (10), the parameters and are constants, and are 0.5 and 0.1, respectively.
2. The medical image segmentation method based on the dual-branch decoder according to claim 1, characterized in that, The step of obtaining the cluster centers obtained at each stage in step 3.4.2 includes: Step 3.4.2.1: Convolution and normalization are performed on pixel features and cluster centers respectively, and pixel and cluster center information are fused by k-means cross-attention; Step 3.4.2.2: The fused pixel and cluster center information are aggregated and optimized by using a multi-head attention module, and the cluster centers are obtained; Step 3.4.2.3: The cluster centers output by the multi-head attention module are updated based on a feed-forward network, and the updated cluster centers of each stage are obtained; The expression of k-means cross-attention operation is: (1) (2) In formulas (1) and (2), Conv () is a convolution layer, BN (·) is a batch normalization operator, KCA () is a k-means cross-attention, is a cluster center of the previous layer, is the output of the lth KCA module. The expression of multi-head attention mechanism operation is: (3) In formula (3), MSA is a multi-head self-attention module; The expression of cluster center update is: (4) In equation (4), FFN is a feedforward neural network module, is the first layer FFN output of the module. 3.The medical image segmentation method based on the dual-branch decoder according to claim 1, characterized in that, The steps of aggregating pixels in each cluster in step 3.5 include: Step 3.5.1: Stack multiple layers of kMaX Transformer modules to form an enhanced decoder; Step 3.5.2: Aggregate pixels in each cluster based on the cross-attention mechanism of the kMaX Transformer module and the updated cluster centers to obtain cluster cluster assignment results; The calculation expression of cross-attention is: (5) In equation (5), C and P are query and pixel features, respectively, Qc , Kp , Vp are linear projection query, key and value, respectively; The expression of aggregation is: (6) (7) In the formulas (6) and (7), B, P, and A are respectively the clustering center, the pixel feature, and the cluster assignment, and .
4. A medical image segmentation system based on a two-branch decoder, applied to the medical image segmentation method based on a two-branch decoder in any one of claims 1-3, characterized in that, Including: A medical image feature acquisition module for extracting image features of the medical image; A medical image resolution recovery module for recovering the resolution of the feature map of the medical image; A medical image segmentation prediction module for obtaining a segmentation prediction result of the medical image; A total segmentation loss acquisition module for obtaining an exchange contrast clustering assignment loss and a multi-scale segmentation loss, and obtaining a total segmentation loss based on the exchange contrast clustering assignment loss and the multi-scale segmentation loss; A medical image automatic segmentation module for optimizing the parameters of the hybrid encoder and the enhanced decoder through the total segmentation loss to complete the automatic segmentation of the medical image.