Semi-supervised medical image segmentation method based on interaction of transformer and CNN

Abstract

The application discloses a semi-supervised medical image segmentation method based on interaction of a Transformer and a CNN, and comprises the following steps: the application introduces the Transformer into a medical image semi-supervised segmentation framework, designs C2T and T2C modules for feature interaction between a Transformer branch and a CNN branch, realizes efficient knowledge sharing between the two branches, improves the ability of the segmentation network to capture detailed features and establish global dependency, and simultaneously adds a feature consistency distribution constraint, uses a feature covariance matrix of the Transformer branch of a Teacher model to constrain a feature covariance matrix of a CNN branch of a Student, and uses a feature covariance matrix of the CNN branch of the Teacher model to constrain a feature covariance matrix of the Transformer, so that the semi-supervised framework after the introduction of the Transformer branch is more stable, and more accurate pseudo labels are generated.

Description

Technical Field

[0001] This invention belongs to the field of medical image analysis, and specifically designs a semi-supervised medical image segmentation method based on the interaction of Transformer and CNN. Background Technology

[0002] In recent years, with the development and popularization of medical imaging technology, medical imaging has provided doctors with important reference information for disease diagnosis. Medical imaging technology utilizes different imaging principles and specialized equipment to non-invasively acquire information about the internal tissue structure of patients, allowing for the assessment of their physiological functions. Medical image segmentation is a crucial step in image-assisted therapy; it involves identifying pixels of diseased organs in medical images and obtaining information such as the location and size of these lesions. This is a technically challenging task in medical image analysis. In real-world applications, the boundaries between different organs and tissues in medical images are often difficult to distinguish, leading to significant uncertainty and subjectivity in the annotation process. Therefore, it is difficult to obtain large amounts of fully annotated, highly accurate medical image datasets. Furthermore, the accessibility and quantity of medical image data are often limited due to patient privacy concerns. Therefore, how to fully utilize the large amount of unannotated data is a pressing issue that needs to be addressed in current medical image segmentation.

[0003] Semi-supervised learning is an important research direction focusing on how to improve model performance by training networks using a small amount of labeled data and a large amount of unlabeled data. Generally speaking, existing semi-supervised medical image segmentation frameworks mainly fall into two categories: CNN-based semi-supervised medical image segmentation methods and Transformer-based semi-supervised medical image segmentation methods. Since CNNs often focus only on extracting local features and ignore contextual information, while Transformers have the ability to establish global dependencies, how to introduce the Transformer structure into semi-supervised learning has become an important research topic. The paper "Semi-Supervised Medical Image Segmentation via CrossTeaching between CNN and Transformer" proposes building a network based on a cross-teaching framework, where predictions generated by the Transformer are used as pseudo-labels for the CNN, and predictions generated by the CNN are used as pseudo-labels for the Transformer, respectively. This approach demonstrates superior performance compared to existing semi-supervised frameworks on datasets such as ACDC. For example, Chinese patent CN114882047 A, published on August 9, 2022, proposes to build a U-shaped segmentation network that combines CNN and Transformer. By adding residual modules to the Transformer structure and using support vector machines at skip connections to further filter and simplify information, the segmentation accuracy of medical images lacking labeled data is greatly improved. Summary of the Invention

[0004] This invention provides a semi-supervised medical image segmentation method based on the interaction between Transformer and CNN.

[0005] The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN proceeds as follows: First, labeled data is used to train the student network and provide initial parameters. Second, the student network parameters are copied to the teacher network, which then provides pseudo-labels for the unlabeled data. Next, data augmentation is performed on the unlabeled data. Labeled data and augmented unlabeled data are simultaneously input into the student network, which optimizes its parameters using the labeled data, extracts features from the unlabeled data, and outputs its predicted distribution. Augmented unlabeled data is then input into the teacher network, which extracts its features and outputs its predicted distribution. Furthermore, a consistency loss is introduced to constrain the unlabeled data features from both the teacher and student networks, and the information entropy of the teacher network's predicted output is calculated. An information entropy threshold is set to improve the confidence of the pseudo-labels, and the semi-supervised loss is calculated using the pseudo-labels and the predicted distributions of the student network. Finally, the student network passes its network parameters to the teacher network using an exponential moving average for further iterative training.

[0006] A semi-supervised medical image segmentation method based on the interaction of Transformer and CNN includes the following steps:

[0007] (1) Construction of teacher-student networks. Both the teacher and student networks consist of a CNN branch and a Transformer branch. C2T (CNN to Transformer) and T2C (Transformer to CNN) modules are introduced at each stage to enable the interaction of feature extraction between the two branches.

[0008] (2) Online training for students. Utilizing labeled data. Train the student network, where N1 is the number of unlabeled data points. Let i be the i-th training sample.

[0009] The student network parameters S(θ) are initialized and then copied to the teacher network T(θ).

[0010] (3) Generate pseudo-labels. Transfer the unlabeled data... The data is input into the teacher network, and the teacher network's predicted output is used as pseudo-labels for the unlabeled data. Where N u The number of unlabeled data. Let i be the i-th unlabeled image.

[0011] (4) Data augmentation. The CutMix data augmentation method is used to augment the unlabeled data.

[0012] (5) Student Network Optimization and Prediction. Labeled data and enhanced unlabeled data are simultaneously input into the student network. The student network updates its parameters using the labeled data and extracts the features F from the unlabeled data. S And output its predicted distribution P S .

[0013] (6) Feature Consistency Constraint. The enhanced unlabeled data is input into the teacher network, and the unlabeled data features F are extracted. T And output its predicted distribution P T The features extracted from the teacher network are constrained to be consistent with those from the student network.

[0014] (7) Pseudo-label supervised training. Calculate the information entropy of the prediction distribution of the teacher model in step 6. A selection threshold is set to filter pseudo-labels, so that only high-confidence pseudo-labels are retained. and utilize Predicted distribution P of supervised student network S .

[0015] (8) Overall loss optimization. When training the above network, the overall loss is optimized and backpropagation is performed to update the network parameters.

[0016] (9) Network parameter transfer. After the network parameters are updated, the student network transfers its network parameters to the teacher network through an exponential moving average for further iterative training.

[0017] (10) Target image segmentation. For a given target domain image, the segmentation model outputs the probability of the class to which each pixel in the target image belongs, and selects the class with the highest probability as the predicted class of that pixel.

[0018] The Transformer and CNN branch interaction process (C2T module) described in step 1 is as follows:

[0019] (i) At each stage, the features learned by the CNN branch are passed to the Transformer branch through the C2T module. Let the feature map output by the i-th stage of the CNN branch in the student network be denoted as . The probability graph predicted by the fully connected layer is in The dimensions are B×D×H×W. The dimensions are B×Nc×H×W, where B is the size of the training batch, H and W are the height and width of the feature map, respectively, D is the number of channels in the feature map, and Nc is the number of classes. The center of each class is obtained by weighting the two.

[0020] (ii) For the i-th stage of the Transformer branch in the student network, its input feature map is: Each layer passes through three fully connected layers W. Q W K W V Calculate Query(Q), Key(K), and Value(V), where Q, K, and V represent three distinct feature maps, and calculate the similarity matrix S between Q and K. Q_K ;

[0021] (iii) Calculate Q and K from step (ii) and the center of each category obtained in step (i), respectively. The similarity is used to obtain the similarity matrix S. Q_cls and similarity matrix S K_cls ;

[0022] (iv) The similarity matrix S obtained through step (iii) Q_cls and similarity matrix S K_cls The information from the class centers of the CNN branch is incorporated and passed to the corresponding layer of the Transformer branch to optimize matrix S.Q_K The resulting similarity matrix is: Multiplying the similarity matrix by the feature map V yields the enhanced feature F. C2T Finally, the features output by the i-th stage of the CNN branch are compared. The summation yields the feature map output by the i-th stage of the Transformer branch.

[0023] The center of each category described in step (i) The calculation is as follows:

[0024]

[0025] in, Let j be the class probability at index j on the segmentation probability map predicted by the CNN branch in the i-th stage of the student network. Let H and W be the feature vector at index j on the feature map corresponding to the i-th stage of the CNN branch of the student network, where H and W are the height and width of the feature map.

[0026] In step (ii), the similarity matrix S between Q and K Q_K The calculation is as follows:

[0027]

[0028]

[0029] Where D is the channel dimension of the feature map. W is the feature map output by the Transformer branch in the (i-1)th stage. Q W K W V There are three fully connected layers. softmax(·) is the activation function.

[0030] In step (iii), Q and K from step (ii) are calculated respectively, along with the centers of each category obtained in step (i). The similarity is used to obtain the similarity matrix S. Q_cls and similarity matrix S K_cls Specifically, it includes:

[0031]

[0032] Where D is the number of channels in the feature map.

[0033] In step (iv), the fused similarity matrix is: The calculations include:

[0034]

[0035] Feature map The final fusion output specifically includes:

[0036]

[0037] in This refers to the feature map output by fusing information from CNN branches in the i-th stage of the Transformer branch of the student network. Let F be the output feature map of the i-th stage of the CNN branch in the student network. C2T The enhanced feature map, where, V is the fused similarity matrix, and V is the feature map.

[0038] In step 1, the Transformer branch interacts with the CNN branch through the T2C module, specifically including:

[0039] (i) At each stage, the student network passes the features learned by the Transformer to the CNN branch through the T2C module. Let the feature map output by the i-th stage of the Transformer branch in the student network be denoted as... in The dimension is B×N×D, where N is the sequence length, and N = H×W. The feature map is flattened and reshaped to obtain the dimension B×D×H×W, where B is the training batch size, D is the number of feature channels, and H and W are the height and width of the feature map, respectively. The segmentation probability map predicted by the fully connected layer is... Its dimensions are B×Nc×H×W, where Nc is the number of categories. The center of each category is obtained by taking a weighted average of the two.

[0040] (ii) The output feature map of the i-th stage of the CNN branch of the student network is Using the Transformer branch and the feature maps from the class centers obtained in step (i) and the CNN branch, the cross-attention mechanism is used to obtain the enhanced feature F. T2C ;

[0041] (iii) The feature F obtained in step (ii) T2C CNN branch features The final output is obtained by fusion.

[0042] In step (i), the center of each category The calculation process is as follows:

[0043]

[0044] in, Let j be the class probability at index position j of the segmentation probability graph predicted by the fully connected layer in the i-th stage of the Transformer branch in the student network. Let J be the feature vector corresponding to the index position j of the feature map output by the i-th stage of the Transformer branch in the student network, where H and W are the height and width of the feature map, respectively.

[0045] In step (ii), the enhanced feature F is obtained by calculating the cross-attention mechanism. T2C Specifically, it includes:

[0046]

[0047]

[0048] Among them W Q W K W V It consists of three fully connected layers. Let D be the feature map output by the i-th stage of the CNN branch in the student network, where D is the number of channels in the feature map, and softmax(·) is the activation function. Let Q, K, and V be the centers of each category calculated in step (i), and let Q, K, and V be the three different feature maps calculated.

[0049] In step (iii), the output features are... The calculations include:

[0050]

[0051] Where F T2C For the enhanced features, This is the feature map output by the i-th stage of the CNN branch in the student network.

[0052] Step 2, which involves initializing the student network and copying its parameters to the teacher network, is as follows:

[0053] (i) Using labeled data Train the student network S(θ) and calculate the supervised training loss using the cross-entropy loss function. Where N l This represents the number of labeled data.

[0054] (ii) After training, its parameters are copied to the teacher network T(θ).

[0055]

[0056] This indicates that the batch size of the input for model training is . There is labeled data, lce Represents the cross-entropy loss function. This represents the true label of the i-th labeled image. Let be the i-th labeled image. The cross-entropy loss function can be expressed as: Where H and W are the height and width of the predicted output segmentation probability map, and y i p represents the true class at position i in the image. i Let i be the predicted class probability at position i.

[0057] The specific process for obtaining pseudo-labels for unlabeled data in step 3 is as follows:

[0058]

[0059] in For the i-th unlabeled image, The pseudo-labels predicted by the teacher network T(θ).

[0060] The specific steps of the data augmentation method CutMix described in step 4 are as follows:

[0061]

[0062] Where M∈{0,1} W×H This is a mask generated to crop out certain regions, where · represents the element-wise dot product, and λ follows a Beta distribution, i.e., λ ~ Beta(α, α). These are the i-th and j-th unlabeled images, respectively. These are the pseudo-labels for the i-th and j-th unlabeled images, respectively.

[0063] Step 5 involves simultaneously inputting labeled data and enhanced unlabeled data into the student network. The student network uses the labeled data to update its network parameters and extracts the features F from the unlabeled data. S And output its predicted distribution P S .

[0064] (i) The student network updates its parameters using labeled data and calculates the supervised training loss using the cross-entropy loss function.

[0065] (ii) Extract unlabeled data features F from the two branches in the student network respectively. S-T and F S-C And output its predicted distribution P S-T and P S-C Ultimately, the mean of the two values ​​is used as the predicted distribution P of the student network. S .

[0066] Supervision of loss in step (i) for:

[0067]

[0068] This indicates that the batch size of the input for model training is . There is labeled data, l ce The cross-entropy loss function can be expressed as: Where H and W are the height and width of the predicted output segmentation probability map, and y m p represents the true class at position m in the image. m Let m be the predicted probability at position m. This represents the predicted segmentation result after inputting the i-th labeled image into the student network. Let represent the label of the i-th labeled image.

[0069] In step (ii), the distribution P is predicted. S The calculation process is as follows:

[0070]

[0071] in and These are the feature extractors for the Transformer and CNN branches of the student network, respectively. and These are the predicted segment heads for the two branches of the student network, respectively. S-T and P S-C For the predicted outputs of the two branches of the student network, x u This is an unlabeled image.

[0072] Step 6: Input the enhanced unlabeled data into the teacher network and extract the unlabeled data features F. T And output its predicted distribution P T The features extracted from the teacher network are constrained to maintain consistency with the student network. The specific steps are as follows:

[0073] (i) Teacher network extracts unlabeled data features F T And output its predicted distribution P T .

[0074] (ii) Extracting features F from the Transformer branch and CNN branch of the teacher network T-T F T-C The features F extracted from the Transformer branch and CNN branch in the student network S-T F S-C Multiplying each of the two by its own transpose yields the covariance feature map.

[0075] (iii) The covariance feature maps of corresponding branches of the teacher network and student network are used to apply consistency constraints using the MSE loss function, and the consistency loss is calculated.

[0076] In step (i), the distribution P is predicted. T The calculation process is as follows:

[0077]

[0078] in and These are the feature extractors for the Transformer and CNN branches of the teacher network, respectively. and These are the segmentation prediction heads for the two branches of the teacher network, respectively. S-T and P S-C For the predicted outputs of the two branches of the teacher network, x u Input image without label.

[0079] The calculation process of the covariance matrix in step (ii) is as follows:

[0080] F Σ =F·(F) T

[0081] Consistency loss in step (iii) The calculation process is as follows:

[0082]

[0083] Where MSE is the squared difference, i.e., MSE(x,y)=x 2 -y 2 .

[0084] Step 7 describes calculating the information entropy for the predicted distribution of the teacher model in step 6. A selection threshold is set to filter pseudo-labels, so that only high-confidence pseudo-labels are retained. and utilize Predicted distribution P of supervised student network S The specific steps are as follows:

[0085] (i) Calculate the information entropy of the predicted distribution of the teacher model.

[0086] (ii) Set a selection threshold to filter pseudo-labels in order to obtain high-confidence prediction output.

[0087] (iii) Utilizing high-confidence labels Predicted distribution P of supervised student network S .

[0088] In step (i), the information entropy of the predicted distribution of the teacher model is calculated:

[0089]

[0090] Where ij is the index of the predicted output, i represents the i-th unlabeled image, and j represents the j-th pixel. C is the number of output channels, and c is the channel index. Let be the pseudo-label for the i-th unlabeled image.

[0091] In step ii, a selection threshold is set to filter pseudo-labels in order to obtain high-confidence prediction outputs. for:

[0092]

[0093] Where γ is the set threshold, The entropy is calculated. Let argmax be the predicted value of the c-th class at the j-th position of the pseudo-label in the i-th unlabeled image. c (·) indicates the category index corresponding to the maximum value. This is the pseudo-label for the j-th pixel of the i-th unlabeled image after filtering.

[0094] Step (iii) utilizes high-confidence labels Predicted distribution P of supervised student network S The specific steps are as follows:

[0095]

[0096] in For unsupervised loss functions, Indicates batch size is unlabeled training data, l ce This is the cross-entropy loss function.

[0097] The overall loss in step 8 is optimized. The monitoring loss in steps 2 and 5 Consistency loss in step 6 and the unsupervised loss in step 7 Linear combination:

[0098]

[0099] Here, λ1 and λ2 are hyperparameters used to balance the effects of the two types of loss on the total loss.

[0100] In step 9, after the network parameters are updated, the student network transmits its network parameters to the teacher network via an exponential moving average (EMA) for further iterative training. The specific process is as follows:

[0101] T(θ)=EMA(S(θ))

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

[0103] (1) Compared with other methods, the semi-supervised medical image segmentation method based on Transformer only interacts with the final prediction result. This method can more fully integrate the local features extracted by CNN with the global features captured by Transformer, thereby improving network performance.

[0104] (2) Compared with other methods, semi-supervised medical image segmentation methods based on Transformer usually calculate pixel-level loss. This method improves the model by calculating the loss of the covariance matrix, which enhances the robustness of the model.

[0105] (3) This method uses two branches in the student network and two branches in the teacher network for cross-teaching at the same time, which can generate more stable and accurate pseudo-labels. Attached Figure Description

[0106] Figure 1 This is a diagram illustrating the overall structure of a semi-supervised medical image segmentation method based on the interaction of Transformer and CNN.

[0107] Figure 2 A diagram of the C2T module structure for the interaction between the Transformer branch and the CNN branch at each stage;

[0108] Figure 3 A diagram of the T2C module structure for the interaction between the Transformer branch and the CNN branch at each stage;

[0109] Figure 4 This is a graph based on the consistency loss constraint of feature distribution;

[0110] Figure 5 This is a flowchart illustrating the semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to the present invention.

[0111] Figure 6 This is a diagram showing the segmentation results of this invention on the Heart Segmentation Dataset (ACDC).

[0112] Figure 7 This is a diagram showing the segmentation results of the present invention on the Skin Disease Segmentation Dataset (ISIC). Detailed Implementation

[0113] The overall framework of the method of this invention is as follows: Figure 1 and Figure 5 As shown, it includes the following steps:

[0114] 1. Construction of Teacher-Student Networks. Both the teacher and student networks consist of a CNN branch and a Transformer branch, with C2T (CNN to Transformer) and T2C (Transformer to CNN) modules introduced at each stage to enable interaction between the two branches in extracting features.

[0115] 2. Online training for students. Utilizing labeled data. Train the student network to initialize the student network parameters S(θ), and then copy these parameters to the teacher network T(θ).

[0116] 3. Generate pseudo-labels. This involves generating pseudo-labels from unlabeled data. The data is input into the teacher network, and the teacher network's predicted output is used as pseudo-labels for the unlabeled data.

[0117] 4. Data Augmentation. The CutMix data augmentation method is used to augment both labeled and unlabeled data.

[0118] 5. Student Network Optimization and Prediction. Augmented labeled and unlabeled data are simultaneously input into the student network. The student network uses the labeled data to update its parameters and extracts the features F from the unlabeled data. S And output its predicted distribution P S .

[0119] 6. Feature Consistency Constraint. The enhanced unlabeled data is input into the teacher network, and the unlabeled data features F are extracted. T And output its predicted distribution P T The features extracted from the teacher network are constrained to be consistent with those from the student network.

[0120] 7. Pseudo-label supervised training. Calculate the information entropy of the prediction distribution of the teacher model from step 6. A selection threshold is set to filter pseudo-labels, so that only high-confidence pseudo-labels are retained. and utilize Predicted distribution P of supervised student network S .

[0121] 8. Overall Loss Optimization. During the training of the above network, the overall loss is optimized and backpropagation is performed to update the network parameters.

[0122] 9. Network Parameter Transmission. After the network parameters are updated, the student network transmits its network parameters to the teacher network via an exponential moving average for further iterative training.

[0123] 10. Target Image Segmentation. For a given target domain image, the segmentation model outputs the probability of the class to which each pixel in the target image belongs, and selects the class with the highest probability as the predicted class of that pixel.

[0124] In step 1, both the teacher and student networks consist of a CNN branch and a Transformer branch. C2T (CNN to Transformer) and T2C (Transformer to CNN) modules are introduced at each stage to enable interaction between the two branches in extracting features.

[0125] Specifically:

[0126] (i)C2T module

[0127] At each stage, the features learned by the CNN branch are passed to the Transformer branch through the C2T module. Let the feature map output by the i-th stage of the CNN branch in the student network be denoted as . The probability graph predicted by the fully connected layer is By taking a weighted average of the two, the center of each category can be obtained.

[0128]

[0129] in, Let j be the class probability at index j on the segmentation probability map predicted by the CNN branch in the i-th stage of the student network. Let H and W be the feature vector at index j on the feature map corresponding to the i-th stage of the CNN branch of the student network, where H and W are the height and width of the feature map.

[0130] For the i-th stage of the Transformer branch in the student network, its input feature map is: Each layer passes through three fully connected layers W. Q W K W V Calculate Query(Q), Key(K), and Value(V), and obtain the similarity matrix S between Q and K. Q_K :

[0131]

[0132]

[0133] Where D is the channel dimension of the feature map. W is the feature map output by the Transformer branch of the student network in the (i-1)th stage. Q W K W V There are three fully connected layers. softmax(·) is the activation function.

[0134] The similarity matrix S is obtained by calculating Q and K in formula (2) and the similarity of the category centers calculated by formula (1). Q_cls and similarity matrix S K_cls .

[0135]

[0136] Where D is the number of channels in the feature map.

[0137] This information is passed to the corresponding layer of the Transformer to optimize matrix S. Q_K The resulting similarity matrix is:

[0138]

[0139] Finally, the features of the i-th stage of the CNN branch are... and the enhanced feature F C2T The fusion yields the final fused output features of the i-th stage of the Transformer branch.

[0140]

[0141] in This refers to the feature map output by fusing information from CNN branches in the i-th stage of the Transformer branch of the student network. Let F be the output feature map of the i-th stage of the CNN branch in the student network. C2T The enhanced feature map, where, V is the fused similarity matrix, and V is the feature map.

[0142] (ii) T2C module

[0143] The student network passes the features learned by the Transformer to the CNN branch through the T2C module at each stage. Let the feature map output by the i-th stage of the Transformer branch in the student network be denoted as . The feature map is flattened and transformed to obtain dimensions B×D×H×W, where B is the training batch size, D is the number of feature channels, and H and W are the height and width of the feature map, respectively. The segmentation probability map predicted by the fully connected layer has dimensions B×Nc×H×W, where Nc is the number of classes. The center of each class can be obtained by weighted averaging of the two values.

[0144]

[0145] in, Let j be the class probability at index position j of the segmentation probability graph predicted by the fully connected layer in the i-th stage of the Transformer branch in the student network. Let J be the feature vector corresponding to the index position j of the feature map output by the i-th stage of the Transformer branch in the student network, where H and W are the height and width of the feature map, respectively.

[0146] The output feature map of the i-th stage of the CNN branch of the student network is The enhanced feature F can be obtained by calculating the cross-attention mechanism using the class center obtained by the Transformer branch through formula (6) and the feature map output by the CNN branch. T2C .

[0147]

[0148]

[0149] Among them W Q W K W V It consists of three fully connected layers. The category center is calculated by formula (6), and D is the number of channels in the feature map.

[0150] The characteristic F obtained from formula (7) T2C The feature map output by the i-th stage of the CNN branch of the student network The final output features are obtained by fusion.

[0151]

[0152] In step 2, labeled data is used. Train the student network S(θ) and calculate the supervised training loss using the cross-entropy loss function. After training, its parameters are copied to the teacher network T(θ), as shown in formula (9): After training, its parameters are copied to the teacher network T(θ).

[0153]

[0154] This indicates that the batch size of the input for model training is . There is labeled data, l ce Represents the cross-entropy loss function. This represents the true label of the i-th labeled image. Let be the i-th labeled image. The cross-entropy loss function can be expressed as: Where H and W are the height and width of the predicted output segmentation probability map, and y i p represents the true class at position i in the image. i Let i be the predicted class probability at position i.

[0155] In step 3, the unlabeled data The data is input into the teacher network, and the teacher network's predicted output is used as pseudo-labels for the unlabeled data. As shown in formula (10):

[0156]

[0157] in For the i-th unlabeled image, The pseudo-labels predicted by the teacher network T(θ).

[0158] In step 4, the CutMix data augmentation method is used to augment the unlabeled data, as shown in formula (11):

[0159]

[0160] Where M∈{0,1} W×H This is a mask generated to crop out certain regions, where · represents the element-wise dot product, and λ follows a Beta distribution, i.e., λ ~ Beta(α, α). These are the i-th and j-th unlabeled images, These are the pseudo-labels for the i-th and j-th unlabeled images, respectively.

[0161] In step 5, both labeled data and enhanced unlabeled data are simultaneously input into the student network. The student network uses the labeled data to update its network parameters and extracts the features F from the unlabeled data. S And output its predicted distribution P S Specifically:

[0162] 5.1 The student network updates network parameters using labeled data and calculates supervised training loss using the cross-entropy loss function. As shown in formula (12):

[0163]

[0164] This indicates that the batch size of the input for model training is . There is labeled data, l ce The cross-entropy loss function can be expressed as: Where H and W are the height and width of the predicted output segmentation probability map, and y m p represents the true class at position m in the image. m Let m be the predicted probability at position m. This represents the predicted segmentation result after inputting the i-th labeled image into the student network. Let represent the label of the i-th labeled image.

[0165] 5.2 Extract unlabeled data features F from the two branches of the student network respectively. S-T and F S-C And output its predicted distribution P S-T and P S-C Ultimately, the mean of the two values ​​is used as the predicted distribution P of the student network. S As shown in formula (13):

[0166]

[0167] in and These are the feature extractors for the Transformer and CNN branches of the student network, respectively. and P represents the classifiers for the two branches of the student network. S-T and P S-C For the predicted outputs of the two branches of the student network, x u This represents the input unlabeled image.

[0168] In step 6, the enhanced unlabeled data is input into the teacher network, and the unlabeled data features F are extracted. T And output its predicted distribution P T The features extracted from the teacher network are constrained to be consistent with those from the student network. Specifically:

[0169] 6.1 Extracting Unlabeled Data Features F from Teacher Networks T And output its predicted distribution P T As shown in formula (14):

[0170]

[0171] in and These are the feature extractors for the Transformer and CNN branches of the teacher network, respectively. and These are the classifiers for the two branches of the teacher network, respectively. T-T and PT-C For the predicted outputs of the two branches of the teacher network, x u This represents the input unlabeled image.

[0172] 6.2 Extracting features F from teacher and student networks T-T F T-C F S-T F S-C Multiplying it by its own transpose yields the covariance feature map. As shown in formula (15), the covariance feature maps of the corresponding branches of the teacher network and student network are constrained using the MSE loss function, and the consistency loss is calculated. To improve the robustness and generalization ability of the model, as shown in Equation (16).

[0173]

[0174]

[0175] Where MSE is the squared difference, i.e., MSE(x,y)=x 2 -y 2 .

[0176] In step 7, the information entropy is calculated for the predicted distribution of the teacher model from step 6. A selection threshold is set to filter pseudo-labels, so that only high-confidence pseudo-labels are retained. and utilize Predicted distribution P of supervised student network S Specifically:

[0177] 7.1 Calculate the information entropy of the prediction distribution of the teacher model, as shown in formula (17); and set a selection threshold to filter pseudo-labels to obtain a high-confidence prediction output, as shown in formula (18):

[0178]

[0179] Where ij is the index of the predicted output, i represents the i-th unlabeled image, and j represents the j-th pixel. C is the number of output channels, and c is the channel index. Let be the pseudo-label for the j-th pixel in the i-th unlabeled image.

[0180]

[0181] Where γ is the set threshold, The entropy is calculated. Let argmax be the predicted value of the c-th class at the j-th position of the pseudo-label in the i-th unlabeled image. c(·) indicates the category index corresponding to the maximum value. This is the pseudo-label for the j-th pixel of the i-th unlabeled image after filtering.

[0182] 7.2 Utilizing high-confidence labels Predicted distribution P of supervised student network S As shown in formula (19):

[0183]

[0184] in For unsupervised loss functions, Indicates batch size is unlabeled training data, l ce This is the cross-entropy loss function.

[0185] Overall loss in step 8 The monitoring loss in steps 2 and 5 Consistency loss in step 6 and the unsupervised loss in step 7 The linear combination of is shown in formula (20):

[0186]

[0187] Here, λ1 and λ2 are hyperparameters used to balance the effects of the two types of loss on the total loss.

[0188] In step 9, after the network parameters are updated, the student network transmits its network parameters to the teacher network via an exponential moving average (EMA) for further iterative training, as shown in formula (21):

[0189] T(θ)=EMA(S(θ)) (21)

[0190] In step 10, the target image segmentation involves, given target domain image data, the segmentation model predicts the class probability for each pixel, and selects the class with the highest predicted probability as the predicted class to obtain the final segmentation mask.

[0191] Furthermore, such as Figure 1 As shown, the specific process of a semi-supervised medical image segmentation method based on the interaction of Transformer and CNN is as follows:

[0192] 1. Teacher-Student Network Construction. Both the teacher and student networks consist of a CNN branch and a Transformer branch. C2T (CNN to Transformer) is introduced at each stage. Figure 2 As shown, the specific steps are as follows:

[0193] (i) At each stage, the features learned by the CNN branch are passed to the Transformer branch through the C2T module. Let the feature map output by the i-th stage of the CNN branch in the student network be denoted as . The probability graph predicted by the fully connected layer is By taking a weighted average of the two, the center of each category can be obtained. As shown below:

[0194]

[0195] in, Let j be the class probability at index j on the segmentation probability map predicted by the CNN branch in the i-th stage of the student network. Let H and W be the feature vector at index j on the feature map corresponding to the i-th stage of the CNN branch of the student network, where H and W are the height and width of the feature map.

[0196] (ii) For the i-th stage of the Transformer branch in the student network, its input feature map is: Each layer passes through three fully connected layers W. Q W K W V Three different feature maps, Query(Q), Key(K), and Value(V), are calculated, and the similarity matrix S between Q and K is also calculated. Q_K .

[0197]

[0198]

[0199] Where D is the channel dimension of the feature map. W is the feature map output by the Transformer branch in the (i-1)th stage. Q W K W V It consists of three fully connected layers.

[0200] (iii) Calculate Q and K from step (ii) and the similarity of the category centers obtained from step (i) to obtain the similarity matrix S. Q_cls and similarity matrix S K_cls .

[0201]

[0202] Where D is the number of channels in the feature map.

[0203] (iv) Pass this information to the appropriate Transformer layer to optimize S Q_KThe resulting similarity matrix is:

[0204]

[0205] (v) Finally, the features from the i-th stage of the CNN branch and the i-th stage of the Transformer branch are fused to obtain...

[0206]

[0207] in This refers to the feature map output by fusing information from CNN branches in the i-th stage of the Transformer branch of the student network. Let F be the output feature map of the i-th stage of the CNN branch in the student network. C2T The enhanced feature map, where, V is the fused similarity matrix, and V is the feature map.

[0208] Similarly, T2C modules are introduced at each stage, such as Figure 3 As shown, the specific steps are as follows:

[0209] (i) At each stage, the student network passes the features learned by the Transformer to the CNN branch through the T2C module. Let the feature map output by the i-th stage of the Transformer branch in the student network be denoted as... The feature map is flattened and transformed to obtain dimensions B×D×H×W, where B is the training batch size, D is the number of feature channels, and H and W are the height and width of the feature map, respectively. The segmentation probability map predicted by the fully connected layer has dimensions B×Nc×H×W, where Nc is the number of classes. The center of each class can be obtained by weighted averaging of the two values.

[0210]

[0211] in, Let j be the class probability at index position j of the segmentation probability graph predicted by the fully connected layer in the i-th stage of the Transformer branch in the student network. Let J be the feature vector corresponding to the index position j of the feature map output by the i-th stage of the Transformer branch in the student network, where H and W are the height and width of the feature map, respectively.

[0212] (ii) The output feature map of the i-th stage of the CNN branch of the student network is The category centers obtained in step (i) can be used through the Transformer branch. The enhanced features obtained by calculating the cross-attention mechanism from the feature maps of the CNN branches are F. T2C .

[0213]

[0214]

[0215] Among them W Q W K W V It consists of three fully connected layers. Let D be the category center calculated in step (i), and D be the number of channels in the feature map.

[0216] (iii) The feature F obtained in step (ii) T2C The final output features are obtained by fusing them with the features from the CNN branches.

[0217]

[0218] 2. Online training for students. Utilizing labeled data. Train the student network to initialize the student network parameters S(θ), and then copy these parameters to the teacher network T(θ).

[0219] 3. Generate pseudo-labels. This involves generating pseudo-labels from unlabeled data. The data is input into the teacher network, and the teacher network's predicted output is used as pseudo-labels for the unlabeled data.

[0220] 4. Data Augmentation. The CutMix data augmentation method is used to augment unlabeled data.

[0221] 5. Student Network Optimization and Prediction. Labeled and enhanced unlabeled data are simultaneously input into the student network. The student network uses the labeled data to update its parameters and extracts the features F from the unlabeled data. S And output its predicted distribution P S .

[0222] 6. Feature Consistency Constraint. The enhanced unlabeled data is input into the teacher network, and the unlabeled data features F are extracted. T And output its predicted distribution P T The features extracted from the teacher network are constrained to be consistent with those from the student network, such as... Figure 4 As shown, the specific process is as follows:

[0223] The features F extracted from the teacher network and student network T-T F T-C F S-T F S-CMultiplying it by its own transpose yields the covariance feature map. The covariance feature maps of corresponding branches in the teacher and student networks are used to apply consistency constraints using the MSE loss function, and the consistency loss is calculated. To improve the robustness and generalization ability of the model, as shown in the following formula.

[0224] F Σ =F·(F) T

[0225]

[0226] Where MSE is the squared difference, i.e., MSE(x,y)=x 2 -y 2 .

[0227] 7. Pseudo-label supervised training. Calculate the information entropy of the prediction distribution of the teacher model from step 6. A selection threshold is set to filter pseudo-labels, so that only high-confidence pseudo-labels are retained. and utilize Predicted distribution P of supervised student network S .

[0228] 8. Overall Loss Optimization. During the training of the above network, the overall loss is optimized and backpropagation is performed to update the network parameters.

[0229] 9. Network Parameter Transmission. After the network parameters are updated, the student network transmits its network parameters to the teacher network via an exponential moving average for further iterative training.

[0230] 10. Target Image Segmentation. For a given target domain image, the segmentation model outputs the probability of the class to which each pixel in the target image belongs, and selects the class corresponding to the highest probability as the predicted class of that point.

[0231] This invention uses two multi-center public datasets (the dermatology segmentation dataset ISIC and the cardiac segmentation dataset ACDC) to evaluate the performance of the present invention. The dermatology segmentation dataset ISIC contains two categories and a total of 2594 images. 1838 images were selected for training, and the remaining 756 images were used for model validation. The cardiac segmentation dataset ACDC contains three categories (left atrium, right atrium, and ventricular wall) and includes images from 100 patients. Images from 70 patients were selected for training, 10 for validation, and 20 for testing. For the semi-supervised task, 3% and 10% of the training sets were selected as labeled data, respectively, with the remainder as unlabeled data. We compared the performance with the current state-of-the-art semi-supervised framework CTCT on both datasets. The ground truth labels were provided by domain experts. The segmentation results are shown below. Figure 6 and Figure 7 As shown, it can be observed that on both datasets, the generalization ability and segmentation performance of the model of this invention are superior to the best-performing existing model structure, verifying the performance of this invention.

Claims

1. A semi-supervised medical image segmentation method based on the interaction of Transformer and CNN, characterized in that, Includes the following steps: (1) Construction of teacher-student network: Both the teacher network and the student network are composed of CNN branch and Transformer branch. The CNN branch extracts medical image features and the Transformer branch extracts medical image features and interacts with each other. The CNN branch interacts with the Transformer branch through the C2T module, and the Transformer branch interacts with the CNN branch through the T2C module. (2) Student network training: The student network is trained using labeled data, the student network parameters are initialized, and the parameters are copied to the teacher network; (3) Generating pseudo-labels: Inputting unlabeled data from medical images into the teacher network, and using the prediction output of the teacher network as pseudo-labels for the unlabeled data; (4) Data augmentation: Using data augmentation methods to augment unlabeled data; (5) Student network optimization and prediction: Labeled data and enhanced unlabeled data are simultaneously input into the student network. The student network uses the labeled data to update network parameters and extracts features from the unlabeled data. And output its predicted distribution ; (6) Feature consistency constraint: The enhanced unlabeled data is input into the teacher network to extract the features of the unlabeled data. And output its predicted distribution Consistency constraints are imposed on the features extracted from the teacher's network and the features extracted from the student's network. (7) Pseudo-label supervised training; prediction distribution of the teacher network in step 6. Calculate information entropy A selection threshold is set to filter pseudo-labels, retaining only high-confidence pseudo-labels. and utilize Predicted distribution of supervised student networks ; (8) Overall Loss Optimization: When training the above network, the overall loss is optimized and backpropagation is performed to update the network parameters: (9) Network parameter transfer: After the network parameters are updated, the student network transfers its network parameters to the teacher network through an exponential moving average for further iterative training; (10) Target image segmentation: For a given target domain image, the segmentation model outputs the probability of the category to which each pixel belongs in the target image, and selects the category with the highest probability as the predicted category of the pixel.

2. The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to claim 1, characterized in that, In step 1, the CNN branch interacts with the Transformer branch through the C2T module, specifically including: (i) At each stage, the features learned by the CNN branch are passed to the Transformer branch through the C2T module, and the CNN branch in the student network is denoted as the th branch. The feature maps output at each stage are The probability graph predicted by the fully connected layer is ,in The dimension is , The dimension is B represents the training batch size, H and W represent the height and width of the feature map, respectively, D represents the number of channels in the feature map, and Nc represents the number of classes. The center of each class is obtained by taking a weighted average of the two values. ; (ii) For the Transformer branch in the student network Each stage has an input feature map as follows: Each layer passes through three fully connected layers. , , Calculate Query(Q), Key(K), and Value(V), where Q, K, and V represent three distinct feature maps, and calculate the similarity matrix between Q and K. ; (iii) Calculate Q and K from step (ii) and the center of each category obtained in step (i), respectively. The similarity is used to obtain the similarity matrix. and similarity matrix ; (iv) Obtained through step (iii) and The information from the class centers of the CNN branch is incorporated and passed to the corresponding layers of the Transformer branch to optimize the matrix. The resulting similarity matrix is: The enhanced features are obtained by multiplying the similarity matrix with the feature map V. Finally, the features output by the i-th stage of the CNN branch are compared. The summation yields the feature map output by the i-th stage of the Transformer branch. .

3. The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to claim 2, characterized in that, In step (i), the center of each category The calculation is as follows: = ； in, Let j be the class probability at index j on the segmentation probability map predicted by the CNN branch in the i-th stage of the student network. Let H and W be the feature vector at index j on the feature map corresponding to the i-th stage of the CNN branch of the student network, where H and W are the height and width of the feature map.

4. The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to claim 2, characterized in that, In step (ii), the similarity matrix between Q and K The calculation is as follows: Q = ,K = , ； Where D is the channel dimension of the feature map. This is the feature map output by the Transformer branch in the (i-1)th stage. , , It consists of three fully connected layers. This is the activation function.

5. The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to claim 2, characterized in that, In step (iii), Q and K from step (ii) are calculated respectively, along with the centers of each category obtained in step (i). The similarity is used to obtain the similarity matrix. and Specifically, it includes: ； Where D is the number of channels in the feature map.

6. The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to claim 2, characterized in that, In step (iv), the fused similarity matrix is: The calculations include: = ( ( )) ； In step (iv), the feature map The final fusion output specifically includes: = V, = + ； in This refers to the feature map output by fusing information from CNN branches in the i-th stage of the Transformer branch of the student network. This is the output feature map of the i-th stage of the CNN branch in the student network. The enhanced feature map, where, V is the fused similarity matrix, and V is the feature map.

7. The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to claim 1, characterized in that, In step 1, the Transformer branch interacts with the CNN branch through the T2C module, specifically including: (i) At each stage, the student network passes the features learned by the Transformer to the CNN branch through the T2C module. Let the T2C module be the T2C module of the student network. The feature maps output at each stage are ,in The dimension is N is the sequence length, and N = H. The feature map is flattened and its shape transformed to obtain the dimension. B represents the training batch size, D represents the number of feature channels, and H and W represent the height and width of the feature map, respectively. The segmentation probability map predicted by the fully connected layer is... , Nc is the number of categories. The center of each category is obtained by taking a weighted average of the two. ; (ii) The output feature map of the i-th stage of the CNN branch of the student network is The enhanced features are obtained by calculating the cross-attention mechanism using the category centers obtained in step i and the feature maps of the CNN branch through the Transformer branch. ; (iii) The features obtained in step (ii) CNN branch features The final output feature map is obtained by fusion. .

8. The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to claim 7, characterized in that, In step (i), the center of each category The calculation process is as follows: = ； in, Let j be the class probability at index position j of the segmentation probability graph predicted by the fully connected layer in the i-th stage of the Transformer branch in the student network. Let J be the feature vector corresponding to the index position j of the feature map output by the i-th stage of the Transformer branch in the student network, where H and W are the height and width of the feature map, respectively. In step (ii), the enhanced features obtained through the cross-attention mechanism are calculated. Specifically, it includes: Q = ,K = ,V = ; V； in , , It consists of three fully connected layers. Let D be the feature map output by the i-th stage of the CNN branch in the student network, and D be the number of channels of the feature map. For activation function, Centered on the category; In step (iii), the feature map output after fusion The calculations include: = + ； in, For the enhanced features, This is the feature map output by the i-th stage of the CNN branch in the student network.

9. The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to claim 1, characterized in that, In step 6, the enhanced unlabeled data is input into the teacher network to extract unlabeled data features. And output its predicted distribution The consistency constraints are applied to the features extracted from the teacher's network and the features extracted from the student's network. The specific steps are as follows: (i) Teachers extract features from unlabeled data online And output its predicted distribution ; (ii) Features extracted from the Transformer and CNN branches in the teacher network , Features extracted from Transformer and CNN branches in student networks , Multiplying each product by its own transpose yields the covariance feature map. , , , ; (iii) The covariance feature maps of corresponding branches of the teacher network and student network are used to apply consistency constraints using the MSE loss function, and the consistency loss is calculated. Through consistency loss Consistency constraints are applied to the features extracted by teachers and students from the network.

10. The semi-supervised medical image segmentation method based on the interaction of Transformer and CNN according to claim 9, characterized in that, In step (iii), consistency loss The calculation process is as follows: ； Where MSE is the calculated squared difference.

Images