Semi-supervised cardiac image segmentation method based on mamba-transformer double structure and contrastive learning

Abstract

The application belongs to the technical field of medical image processing, and particularly relates to a semi-supervised cardiac image segmentation method based on Mamba-Transformer double structure and contrast learning. The segmentation model used by the method comprises an MTSeg encoder and a VNet decoder. After linear layer and position embedding operation, an input image enters the MTSeg encoder. The MTSeg encoder comprises a plurality of MTSeg encoding modules. The VNet decoder comprises a plurality of decoding modules. The output feature of the MTSeg encoding module is connected with the corresponding decoding module through a projection operation. The output feature vector of the VNet decoder is subjected to a normalization operation to obtain a segmentation result. The MTSeg encoding module comprises an improved Mamba branch and a Transformer branch. The output feature vectors of the two branches are fused through a cross attention layer to obtain the output feature vector of the MTSeg encoding module. The improved Mamba branch comprises a series connection of an MS module and a DRFB module. The MS module is used to enhance the feature extraction and calculation efficiency of the model. The DRFB module is used to enhance the expression ability and training stability of the model. Long-range dependency and global information, multi-scale and local information are effectively integrated, the feature extraction capability of the model is improved, and a more fine and accurate segmentation result is obtained.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of medical image processing, and specifically relates to a semi-supervised cardiac image segmentation method based on Mamba-Transformer double structure and contrast learning. BACKGROUND

[0002] Cardiovascular diseases have become one of the major diseases threatening human health, and cardiac function assessment plays a crucial role in early diagnosis, treatment plan development and efficacy monitoring of cardiovascular diseases, and the structure and morphology of the heart are important basis for assessing the cardiac function. Cardiac CT images provide effective cardiac physiological characteristics and anatomical information, however, the complex anatomical structure of the heart, morphological changes caused by different diseases, and image noise and high-similarity tissues (such as the lung and blood vessels) all increase the difficulty of cardiac segmentation tasks.

[0003] Traditional image segmentation networks, such as convolutional neural networks (CNN) and ViT (Vision Transformer), have achieved certain results in medical image segmentation, but still cannot meet the requirements of high precision and robustness for clinical applications. In addition, high-quality medical image annotation requires experienced experts, and the annotation process is time-consuming and costly, so obtaining large-scale high-quality labeled data has become a bottleneck. SUMMARY

[0004] In view of the deficiencies of the prior art, the technical problem to be solved by the present application is to provide a semi-supervised cardiac image segmentation method based on Mamba-Transformer double structure and contrast learning.

[0005] The technical problem of the present application is solved by adopting the following technical solution:

[0006] A semi-supervised cardiac image segmentation method based on Mamba-Transformer double structure and contrast learning, comprising the following steps:

[0007] Step 1: Obtain a data set and perform preprocessing;

[0008] Step 2: Construct a teacher model and a student model, both of which include an MTSeg encoder and a VNet decoder, the input image is input into the MTSeg encoder after linear layer and position embedding operation, the MTSeg encoder includes multiple MTSeg encoding modules, the VNet decoder includes multiple decoding modules, the output features of the MTSeg encoding module are connected to the corresponding decoding module through projection operation; the output feature vector of the VNet decoder is normalized to obtain a segmentation result;

[0009] The MTSeg encoding module comprises an improved Mamba branch and a Transformer branch, and the output feature vectors of the two branches are fused through a cross-attention layer to obtain the output feature vector of the MTSeg encoding module; the improved Mamba branch comprises a MS module and a DRFB module connected in series, in the MS module, the input feature vector is input into two branches after passing through a linear layer, in one branch, the input feature vector sequentially passes through a linear layer, an FFT module and a SiLU activation layer, in the other branch, the input feature vector sequentially passes through a linear layer, an ACFF module, an SS2D module and a linear layer, the output feature vectors of the two branches are multiplied element by element, and then added element by element with the input feature vector of the MS module, and the feature vector obtained by the addition sequentially passes through a layer normalization, a multi-layer perception and a ReLu activation layer to obtain the output feature vector of the MS module;

[0010] In the FFT module, the input feature vector is converted to the frequency domain through Fourier transform to obtain a frequency domain feature vector; a frequency selection mask is generated, and the frequency domain feature vector is multiplied by the frequency selection mask to obtain a low frequency feature vector, and the remaining part of the frequency domain feature vector is a high frequency feature vector; the low frequency feature vector and the high frequency feature vector are weighted and fused, and then subjected to activation processing to obtain the output feature vector of the FFT module;

[0011] In the ACFF module, the input feature vector is input into two branches, in one branch, the input feature vector sequentially passes through a convolution layer, a fully connected layer and a GELU activation layer, in the other branch, the input feature vector sequentially passes through a convolution layer and a SiLu activation layer, and the output feature vectors of the two branches are added element by element to obtain the output feature vector of the module;

[0012] In the DRFB module, the input feature vector is input into two branches, in one branch, convolution operation is performed, and in the other branch, the input feature vector is multiplied element by element with a filter weight vector generated by a learnable filter, the output feature vectors of the two branches are added element by element, and then sequentially pass through layer normalization and a ReLu activation function to obtain the output feature vector of the module;

[0013] Step three, jointly training the teacher model and the student model by using the memory bank mechanism, taking the trained student model as a segmentation model for segmenting three-dimensional heart images.

[0014] Further, the MTSeg encoder comprises eight MTSeg encoding modules, and the VNet decoder comprises four decoding modules, and the output features of the second, fourth, sixth and eighth MTSeg encoding modules are projected and connected to the corresponding decoding modules.

[0015] Further, for the Transformer branch, the input feature vector is connected with itself in residual connection after normalization and multi-head attention layer, to obtain feature A; the feature A is connected with itself in residual connection after normalization and multi-layer perceptron, to obtain the output feature vector of the Transformer branch.

[0016] Further, the decoding module comprises two up-sampling layers and a 3D convolution layer connected in sequence.

[0017] Further, the model training process is divided into three stages:

[0018] The first stage: using the student model to perform supervised learning on the labeled data and calculating the supervised loss; for unlabeled data, using the teacher model to generate pseudo labels and calculating the unsupervised loss;

[0019] The second stage: using the teacher model to extract feature vectors from the labeled data, screening the feature vectors, and storing the feature vectors with the required prediction accuracy in the memory bank;

[0020] The third stage: using the feature vectors in the memory bank for contrastive learning, i.e. inputting the training set into the student model for feature extraction, and the output feature vector of the student model decoder is processed through two multi-layer perceptrons to obtain the feature vector set P, the process is represented as:

[0021] P=g θ (q θ (f θ -(x))) (2)

[0022] In the formula, x represents the training set, f θ -(x) represents the output feature vector of the student model decoder, g θ , q θ represents a multi-layer perceptron;

[0023] The feature vectors in the feature vector set P are divided into each category subspace according to the category to obtain the feature vector set P1,...,P c ,...,P J , P c ={p c} represents the feature vector set of category c, p c represents the feature vector of category c, and J represents the number of categories; let Z c ={z c} represent the feature vector set of category c stored in the memory bank, z c represents the feature vector of category c stored in the memory bank; the cosine similarity between the feature vector p c and z c is calculated according to the following formula:

[0024]

[0025] The eigenvector p is calculated according to the following formula. c With z c Distance between:

[0026]

[0027] In the formula, Represents the eigenvector p c and z c The attention weights are calculated in the same way for both, where Calculated by the following formula:

[0028]

[0029] In the formula, Let set P be... c The number of feature vectors, S c,θ (·) represents the attention mechanism, p i Describe set P c The feature vector formed by the i-th dimension of all feature vectors in the dataset;

[0030] The comparative loss is calculated using the following formula:

[0031]

[0032] In the formula, Represents set Z c The number of eigenvectors.

[0033] Compared with the prior art, the beneficial effects of the present invention are as follows:

[0034] 1. The MTseg encoder combines the advantages of Transformer networks in capturing long-range dependencies and global information with the powerful capabilities of Mamba networks in multi-scale and local information modeling. This achieves effective integration of long-range dependencies, global information, and multi-scale and local information, enhancing the model's feature extraction capabilities for complex images and resulting in more refined and accurate segmentation results. By introducing a cross-attention mechanism, the model can effectively fuse feature information from different branches, using features from one branch as queries and features from another branch as keys and values. This enhances the interaction between local and global information, improving the accuracy of feature fusion and the overall performance of the model. The MTseg encoding module, through cross-modal feature fusion, enables the transfer and integration of information between different feature spaces, further improving the model's ability to capture complex features and enhancing its performance in segmentation tasks.

[0035] The improved Mamba branch fuses the feature vectors of different frequency components in the frequency domain through the FFT module, optimizes the feature representation, enhances the robustness of the model to noise and the ability to capture detailed information. The ACFF module enhances the feature expression and computational efficiency of the model through multi-path feature fusion and non-linear activation, especially suitable for processing high-dimensional data.

[0036] 2、The VNet decoder has strong 3D data modeling capability, and adopts three-dimensional convolution to more efficiently process the spatial information of three-dimensional heart images. Compared with the traditional U-Net, the VNet better maintains the spatial consistency between voxels, and is more suitable for three-dimensional medical image segmentation tasks.

[0037] 3、In order to reduce the dependence on a large amount of labeled data, the memory bank mechanism is introduced, and the high-confidence features extracted by the teacher model are stored in the memory bank. The contrast learning strategy can better utilize labeled and unlabeled data, significantly improve segmentation accuracy, accelerate model convergence, and improve intra-class consistency and inter-class discriminability. BRIEF DESCRIPTION OF DRAWINGS

[0038] Fig. 1 The structure diagram of the student model and the teacher model;

[0039] Fig. 2 The framework diagram of semi-supervised training. DETAILED DESCRIPTION

[0040] The specific embodiments will be described below in conjunction with the accompanying drawings, which are only used to specifically introduce the technical schemes of the present application, and do not limit the protection scope of the present application.

[0041] The present application provides a semi-supervised cardiac image segmentation method based on Mamba-Transformer double structure and contrast learning (referred to as method, see Figs. 1-2 ), comprising the following steps:

[0042] Step one, obtain a plurality of three-dimensional heart images as a data set, and pre-process the data set, including adjusting the resolution and data enhancement;

[0043] Step two, construct the teacher model and the student model, both of which adopt the same network structure, and both include MTSeg encoder and VNet decoder, the input image is subjected to linear layer and position embedding operation to generate a feature vector, which is input into the MTSeg encoder for encoding; the MTSeg encoder includes a plurality of MTSeg encoding modules, and the VNet decoder includes a plurality of decoding modules, the output features of the MTSeg encoding modules are subjected to projection operation and then connected with the corresponding decoding modules by jumping; this embodiment includes 8 MTSeg encoding modules and 4 decoding modules, the output features of the second, fourth, sixth and eighth MTSeg encoding modules are subjected to projection operation and then connected with the corresponding decoding modules by jumping; the output feature vector of the VNet decoder is subjected to normalization operation to obtain a segmentation result;

[0044] The MTSeg encoding module includes an improved Mamba branch and a Transformer branch, forming a Mamba-Transformer double-branch structure, the output feature vectors of the improved Mamba branch and the Transformer branch are fused by a cross-attention layer to obtain the output feature vector of the MTSeg encoding module; the MTSeg encoding module realizes effective fusion of global information and local detail information, and overcomes the shortcomings of single branch in modeling long-range dependency and capturing local detail information. The improved Mamba branch includes a series of MS modules and a dynamic residual filtering module (Dynamic Residual Filtering Block, DRFB), the MS module is used to enhance the feature extraction and computational efficiency of the model, and the dynamic residual filtering module DRFB is used to enhance the expression ability and training stability of the model. In the MS module, the input feature vector is subjected to linear layer and then enters two branches, in one branch, it is subjected to linear layer, feature optimization module (FFT) and SiLU activation layer in sequence, in the other branch, it is subjected to linear layer, adaptive convolution feature fusion module (ACFF), SS2D module and linear layer in sequence, the output feature vectors of the two branches are multiplied element by element, and then added with the input feature vector of the MS module element by element, the feature vector obtained by addition is subjected to layer normalization, multi-layer perception MLP and ReLu activation layer in sequence to obtain the output feature vector of the MS module.

[0045] In the feature optimization module FFT, the input feature vector is converted to the frequency domain through Fourier transform to obtain a frequency domain feature vector; a frequency selection mask is generated, and the frequency domain feature vector is multiplied by the frequency selection mask to obtain a low-frequency feature vector, and the remaining part of the frequency domain feature vector is a high-frequency feature vector; after the low-frequency feature vector and the high-frequency feature vector are weighted and fused, and then subjected to activation processing, an output feature vector of the feature optimization module is obtained. The low-frequency feature vector usually contains the main structural information of the image, and the high-frequency feature vector contains the detail and edge information, so the fusion of the high-frequency and low-frequency feature vectors helps to capture the detail information in the image, optimizes the representation of the input data, enables the model to better extract and learn features, and thus improves the segmentation accuracy.

[0046] In the ACFF module, the input feature vector enters two branches, in one branch, sequentially passes through a convolution layer, a full connection layer and a GELU activation layer, and in the other branch, sequentially passes through a convolution layer and a SiLu activation layer, the output feature vectors of the two branches are element-wise added to obtain the output feature vector of the module.

[0047] In the DRFB module, the input feature vector enters two branches, in one branch, a convolution operation is performed, and in the other branch, the output feature vector is element-wise multiplied by a filter weight vector generated by a learnable filter, and the output feature vectors of the two branches are element-wise added, and then sequentially pass through layer normalization and a ReLu activation function to obtain the output feature vector of the module. The initial value of the filter weight vector is 1.

[0048] For the Transformer branch, the input feature vector is normalized and passed through a multi-head attention layer, and then is residual connected with itself to obtain feature A; the feature A is normalized and passed through a multi-layer perceptron, and then is residual connected with itself to obtain the output feature vector of the Transformer branch.

[0049] A cross-attention mechanism is introduced in the MTSeg encoding module, aiming to better fuse the output feature vectors of the improved Mamba branch and the Transformer branch. The cross-attention mechanism uses a multi-head attention mechanism, and by taking the output feature vector of one branch as a query and the output feature vector of the other branch as a key and a value, the interaction between local information and global information is effectively enhanced.

[0050] The decoding module includes two up-sampling layers and a 3D convolution layer connected in sequence; unlike the traditional U-Net, the VNet decoder introduces 3D convolution operation, which can fully utilize the spatial information of three-dimensional images.

[0051] Step three, jointly train the teacher model and the student model by using the memory bank mechanism, and use the trained student model as the segmentation model for three-dimensional cardiac image segmentation.

[0052] The model training process is divided into three stages: in the first stage, the student model is supervised to learn the labeled data, and the supervised loss L label is calculated unlabel The teacher model generates pseudo labels for unlabeled data and calculates the unsupervised loss L unlabel , which is used to guide the optimization of the student model; the supervised loss L label and the unsupervised loss L unlabel Both use the combination of dice loss and cross-entropy loss, which is calculated by the following formula in the voxel manner:

[0053]

[0054] In the formula, I represents the number of voxels, J represents the number of classes, Y i,j and G i,j represent the probability of the jth class at voxel i and the one-hot encoding true value;

[0055] The second stage: the teacher model extracts feature vectors from the labeled data, and selects the feature vectors according to the prediction accuracy, and stores the feature vectors that meet the requirements in the memory bank. The memory bank is used to store feature vectors of different classes, combined with the contrast learning strategy, the memory bank narrows the distance between feature vectors of the same class and increases the distance between feature vectors of different classes, effectively enhancing the discrimination ability of the features and significantly improving the segmentation accuracy of the model;

[0056] The third stage: based on the memory bank, the consistency of intra-class features and the separation of inter-class features are optimized; the training set is input into the student model for feature extraction, and the output feature vector of the student model decoder is passed through two multilayer perceptrons to obtain the feature vector set P, the process is represented as:

[0057] P=g θ (q θ (f θ -(x))) (2)

[0058] In the formula, x represents the training set, f θ -(x) represents the output feature vector of the student model decoder, g θ and q θ represent multilayer perceptrons;

[0059] The feature vectors in the feature vector set P are divided into each class subspace according to the class to obtain the feature vector sets P1,...,P c ,...,P J , Pc c denotes the set of feature vectors of class c, p c denotes a feature vector of class c; Z c c denotes the set of feature vectors of class c stored in the memory bank, z c denotes a feature vector of class c stored in the memory bank; the feature vector p c is calculated according to the following formula: c cosine similarity between p

[0060]

[0061] The feature vector p c is calculated according to the following formula: c distance between p

[0062]

[0063] In the formula, denotes the attention weight of the feature vector p c and z c , and the calculation method is the same, where is calculated by the following formula:

[0064]

[0065] In the formula, denotes the number of feature vectors in the set P c , S c,θ (·) denotes the attention mechanism, p i denotes the feature vector formed by the i-th dimension of all feature vectors in the set P c .

[0066] The contrastive loss L contr is calculated according to the following formula:

[0067]

[0068] In the formula, denotes the number of feature vectors in the set Z c .

[0069] The discriminative ability of the model is improved by the contrastive loss L contr , and the student-teacher model realizes collaborative learning on labeled and unlabeled data, fully utilizes unlabeled data, and improves the performance and generalization ability of the model.

[0070] ​​Based on the MMWHS dataset, the ratio of labeled data and unlabeled data is 10% and 90%, and the Dice similarity coefficient (DSC) is used as the evaluation index. Compared with the existing mainstream model, the segmentation results are shown in Table 1.

[0071] Table 1 Segmentation results of different models on the MMWHS dataset

[0072]

[0073] As shown in Table 1, the segmentation model of the present application shows significant advantages in the segmentation task of five parts of the heart (both ventricles, both atria and myocardium). Compared with the existing mainstream model, the average DSC of the five parts is improved by 5.86%, 9.03% and 0.6% respectively, which verifies the effectiveness of the segmentation model. As can be seen, the MTSeg encoder encoder with a double-branch structure can effectively improve the segmentation accuracy of the model, especially the improved Mamba branch can capture local information better, and the cross-attention mechanism is used to complement the global attention of the Transformer branch, enhancing the feature extraction capability of the model. The method of contrast learning can better extract the features of the cardiac image, and enhance the segmentation ability of the model.

[0074] The sources of the above existing mainstream models are as follows:

[0075] 1. Luo X, Liao W, Chen J, et al. Efficient semi-supervised gross target volume of nasopharyngeal carcinoma segmentation via uncertainty rectified pyramid consistency [C] / / Medical Image Computing and Computer Assisted Intervention - MICCAI 2021: 24th International Conference, Strasbourg, France, September 27-October 1, 2021, Proceedings, Part II 24. Springer International Publishing, 2021: 318-329.

[0076] 2. Wang X, Wu Z, Lian L, et al. Debiased learning from naturally imbalanced pseudo-labels [C] / / Proceedings of the IEEE / CVF Conference on Computer Vision and Pattern Recognition. 2022: 14647-14657.

[0077] 3. Wang H, Li X. Dhc: Dual-debiased heterogeneous co-training framework for class-imbalanced semi-supervised medical image segmentation [C] / / International conference on medical image computing and computer-assisted intervention. Cham: Springer Nature Switzerland, 2023: 582-591.

[0078] The unmentioned parts of the present application apply to the prior art.

Claims

1. A semi-supervised cardiac image segmentation method based on Mamba-Transformer dual-structure and contrastive learning, characterized in that, Includes the following steps: Step 1: Obtain the dataset and preprocess it; Step 2: Construct teacher and student models, both including MTSeg encoders and VNet decoders. The input image is processed through linear layers and position embedding operations before entering the MTSeg encoder, which contains multiple MTSeg encoding modules. The VNet decoder contains multiple decoding modules. The output features of the MTSeg encoding modules are projected and then skipped between the corresponding decoding modules. The output feature vector of the VNet decoder is normalized to obtain the segmentation result. The MTSeg encoding module includes an improved Mamba branch and a Transformer branch. The output feature vectors of the two branches are fused through a cross-attention layer to obtain the output feature vector of the MTSeg encoding module. The improved Mamba branch includes a cascaded MS module and a DRFB module. In the MS module, the input feature vector passes through a linear layer and then enters two branches. In one branch, it passes through a linear layer, an FFT module, and a SiLU activation layer in sequence. In the other branch, it passes through a linear layer, an ACFF module, an SS2D module, and a linear layer in sequence. The output feature vectors of the two branches are multiplied element-wise and then added element-wise to the input feature vector of the MS module. The resulting feature vector passes through a layer normalization, a multilayer perceptron, and a ReLU activation layer in sequence to obtain the output feature vector of the MS module. In the FFT module, the input feature vector is transformed to the frequency domain through Fourier transform to obtain the frequency domain feature vector; a frequency selection mask is generated, and the frequency domain feature vector is multiplied by the frequency selection mask to obtain the low-frequency feature vector, and the remaining part of the frequency domain feature vector is the high-frequency feature vector; the low-frequency feature vector and the high-frequency feature vector are weighted and fused, and then activated to obtain the output feature vector of the FFT module. In the ACFF module, the input feature vector enters two branches. In one branch, it passes through a convolutional layer, a fully connected layer, and a GELU activation layer in sequence. In the other branch, it passes through a convolutional layer and a SiLu activation layer in sequence. The output feature vectors of the two branches are added element by element to obtain the output feature vector of the module. In the DRFB module, the input feature vector is fed into two branches. A convolution operation is performed in one branch, and in the other branch, it is multiplied element-wise with the filter weight vector generated by the learnable filter. The output feature vectors of the two branches are then summed element-wise and passed through layer normalization and ReLU activation function in sequence to obtain the output feature vector of the module. Step 3: Use the memory bank mechanism to jointly train the teacher model and the student model, and use the trained student model as the segmentation model for the segmentation of 3D heart images.

2. The semi-supervised cardiac image segmentation method based on Mamba-Transformer dual-structure and contrastive learning according to claim 1, characterized in that, The MTSeg encoder contains 8 MTSeg encoding modules, and the VNet decoder contains 4 decoding modules. The output features of the second, fourth, sixth, and eighth MTSeg encoding modules are projected and then connected to the corresponding decoding modules in a skip connection.

3. The semi-supervised cardiac image segmentation method based on Mamba-Transformer dual-structure and contrastive learning according to claim 1 or 2, characterized in that, For the Transformer branch, the input feature vector is normalized and then passed through a multi-head attention layer, and then residually connected to itself to obtain feature A; feature A is normalized and then passed through a multilayer perceptron, and then residually connected to itself to obtain the output feature vector of the Transformer branch.

4. The semi-supervised cardiac image segmentation method based on Mamba-Transformer dual structure and contrastive learning according to claim 3, characterized in that, The decoding module comprises two upsampling layers and a 3D convolutional layer connected in sequence.

5. The semi-supervised cardiac image segmentation method based on Mamba-Transformer dual-structure and contrastive learning according to claim 1, characterized in that, The model training process is divided into three stages: Phase 1: Supervised learning of labeled data using the student model and calculation of supervised loss; For unlabeled data, pseudo-labels are generated using the teacher model and calculation of unsupervised loss. The second stage involves using a teacher model to extract feature vectors from labeled data, filtering the feature vectors, and storing the feature vectors that meet the prediction accuracy requirements into a memory bank. The third stage: comparative learning using feature vectors from the memory bank, i.e., inputting the training set into the student model for feature extraction, and the output feature vectors of the student model decoder are passed through two multilayer perceptrons to obtain the feature vector set P. The process is represented as follows: P=g θ (q θ (f θ -(x))) (2) In the formula, x represents the training set, and f θ -(x) represents the output feature vector of the student model decoder, g θ q θ Represents a multilayer perceptron; The feature vectors in the feature vector set P are divided into category subspaces according to their categories, resulting in feature vector sets P1,...,P for each category. c ,...,P J P c ={p c } represents the set of feature vectors for category c, p c Let Z represent the feature vector of category c, and J represent the number of categories; let Z c ={z c } represents the set of feature vectors of category c stored in the memory bank, z c This represents the feature vector of category c stored in the memory bank; the feature vector p is calculated according to the following formula. c With z c Cosine similarity between them: The eigenvector p is calculated according to the following formula. c With z c Distance between: In the formula, Represents the eigenvector p c and z c The attention weights are calculated in the same way for both, where Calculated by the following formula: In the formula, Let set P be... c The number of feature vectors, S c, θ(·) represents the attention mechanism, p i Let set P be... c The feature vector formed by the i-th dimension of all feature vectors in the dataset; The comparative loss is calculated using the following formula: In the formula, Represents set Z c The number of eigenvectors.

Images