Semi-supervised medical image decoupled contrast segmentation method based on uncertainty guidance

Abstract

The application discloses a semi-supervised medical image decoupling contrast segmentation method based on uncertainty guidance, relates to the technical field of computer vision and semi-supervised learning, and comprises the following steps: constructing a double student-single teacher same architecture network; selecting a better "winner" student model by predicting labeled data with a Dice loss; updating the teacher model parameters by using an EMA formula; generating pseudo labels for unlabeled data by the teacher model, and guiding the student model to update by the Dice loss; introducing a CCB module, calculating uncertainty according to an entropy value, combining prediction consistency to determine pseudo labels to assist complementary learning; calculating a loss according to a total loss formula, and iteratively training until convergence; and the scheme can improve segmentation accuracy and stability compared with the prior art, can assist in segmenting abnormal regions and lesions of a magnetic resonance image, can reduce the burden of doctors, and is universal and expandable in architecture. Therefore, the problems of model parameter coupling, incorrect prediction misguiding and high labeling dependence of the prior art are effectively solved.

Description

Uncertainty-Guided Semi-Supervised Medical Image Decoupling and Contrast Segmentation Method Technical Field

[0001] This invention relates to the fields of computer vision and semi-supervised learning technology, specifically to a semi-supervised medical image decoupling and contrast segmentation method based on uncertainty guidance. Background Technology

[0002] With the rapid development of computing power and neural networks, a large number of supervised deep learning networks have emerged. Accurate medical image segmentation can provide doctors with key visual information for diagnosis and treatment. However, fully supervised deep learning methods require a large number of medical images annotated by professional doctors. The annotation process is not only time-consuming, labor-intensive, and costly, but also, for 3D medical images such as MRI, experts need to annotate each image in the 3D data, further increasing the annotation burden.

[0003] Semi-supervised medical image segmentation methods can alleviate this problem. They combine training with a small amount of labeled data and a large amount of unlabeled data, reducing reliance on labeled data. Therefore, their application in this field is becoming increasingly widespread. Current mainstream methods are primarily pseudo-labeling and consistency regularization. Pseudo-labeling first trains an initial model using a small amount of labeled data, then uses the model to predict unlabeled data, selecting high-confidence results as "pseudo-labels" to supplement the training set. Consistency regularization requires that different perturbations applied to the same input, or input into different initialization models, should result in consistent output.

[0004] The "Meanteacher" method combines two approaches, containing a student model and a teacher model with identical structures. Student model parameters are updated via gradient descent, while teacher model parameters are updated using the moving exponential average (EMA) of the student model parameters. However, this approach can lead to the two models' predictions becoming increasingly similar in the later stages, resulting in severe parameter coupling. If the student model makes a prediction error, the teacher model can accumulate errors and, in turn, mislead the student model, causing "confirmation bias."

[0005] To address the parameter coupling problem, one approach is to replace the teacher model with another student model, with each student model updating its own parameters. While this can alleviate coupling and increase prediction diversity, there is still a problem where a wrong prediction by one model can mislead the other, ultimately causing the model training to fail.

[0006] In summary, a semi-supervised medical image decoupling and contrast segmentation method based on uncertainty guidance is proposed to address the problem caused by the tight coupling between student model weights and teacher model weights. Summary of the Invention

[0007] In view of this, the technical problem to be solved by this invention is to propose a semi-supervised medical image decoupled contrastive segmentation method based on uncertainty guidance. This method uses a single-teacher, dual-student model. After each iteration, the better-performing student model updates its parameter weights to the teacher model using a moving exponential average (EMA). The teacher model then provides pseudo-labels to the underperforming student model to promote model improvement. Through this dynamic competitive update mechanism, both student models have the opportunity to contribute to the teacher model's update via the moving exponential average (EMA), which naturally reduces the weight coupling between the single student and teacher models. To address the problem of one student model's incorrect prediction leading to incorrect guidance from the other, this paper proposes a dual-student feature contrast and complementarity module. The two student models exchange more certain predictions to achieve more stable training.

[0008] To achieve the above objectives, the present invention provides the following technical solution: a semi-supervised medical image decoupling and contrast segmentation method based on uncertainty guidance, comprising the following steps:

[0009] S1, construct a network architecture consisting of two student models and one teacher model, with all three models having the same network architecture;

[0010] S2, Processing labeled data: Input labeled medical images into two student models, use the Dice loss function to determine the "winner" student model with better predictions, and update the teacher model parameters using the Moving Exponential Mean (EMA) formula; the Moving Exponential Mean (EMA) formula is:

[0011]

[0012] In the formula, These are the parameters of the teacher model in the current iteration. These are the teacher model parameters from the previous iteration. For the "winner" student model parameters;

[0013] S3, Processing unlabeled images: Input the unlabeled medical image into the teacher model to obtain the prediction result of the teacher model; input the prediction result into the softmax activation function to obtain the class probability value, and then determine the class index through the max function, which serves as the pseudo-label of the teacher model; calculate the difference between the output results of the two student models on the unlabeled medical image and the pseudo-label through the Dice loss function, and guide the two student models to update their parameters.

[0014] S4 introduces a complementary module (CCB) to calculate the prediction uncertainty of the two student models based on the entropy value, and determines pseudo-labels based on prediction consistency and uncertainty to guide learning;

[0015] S5. Calculate the loss according to the total loss formula, and iteratively train the network until convergence.

[0016] As a preferred embodiment, the Dice loss function calculation formula described in S2 is as follows:

[0017]

[0018] In the formula, X refers to the model's output predicted value L. S1 ,L S2 Y refers to the true label corresponding to the model output; the two student models are student model 1 and student model 2, and the Dice loss formula for student model 1 is:

[0019]

[0020] In the formula, L S1 Y represents the predicted output probability value of Student Model 1; Y represents the true label value corresponding to the labeled image.

[0021] The Dice loss formula for student model 2 is as follows:

[0022]

[0023] In the formula, L S2 Y represents the predicted output probability value of Student Model 2; Y represents the true label value corresponding to the labeled image.

[0024] As a preferred option, the formula for calculating the loss of student model 1 and the pseudo-label in S3 is:

[0025] =

[0026] In the formula, This represents the output of student model 1 on unlabeled medical images; T represents the pseudo-labels generated by the teacher model.

[0027] The formula for calculating the loss of student model 2 and the pseudo-label loss is as follows:

[0028] =

[0029] In the formula, T represents the output of student model 2 on unlabeled medical images; T represents the pseudo-labels generated by the teacher model.

[0030] In S4, the formula for calculating the mutual learning loss of student models 1 is:

[0031]

[0032] The formula for calculating the mutual learning loss of student model 2 is:

[0033]

[0034] In the formula, "mask" refers to the fact that the two models only calculate the prediction results with low uncertainty, and the remaining samples are not included in this calculation.

[0035] As a preferred embodiment, the formula for calculating the prediction uncertainty of the two-student model based on entropy value described in S4 is as follows:

[0036]

[0037] In the formula, P represents the probability value of the model prediction result, and i represents the number of categories.

[0038] As a preferred embodiment, the loss formula in S5, calculated according to the total loss formula, is as follows:

[0039] Loss= + )

[0040] In the formula, Represented as weighting coefficients, It is calculated based on the prediction results of labeled images and their corresponding labels. It is calculated using the prediction results of the student model and the pseudo-labels generated by the two student models. It is calculated using the prediction results from the student model and the pseudo-labels generated by the teacher model.

[0041] The application of the uncertainty-guided semi-supervised medical image decoupling and contrast segmentation method allows users to accurately segment abnormal regions in magnetic resonance images, assisting doctors in tasks such as lesion segmentation and reducing the burden on doctors in disease diagnosis and prognosis assessment.

[0042] Compared with existing technologies, the uncertainty-guided semi-supervised medical image decoupling and contrast segmentation method provided by this invention has the following advantages:

[0043] (1) Innovation of the "winner" dual-model architecture: This invention designs a new "winner" dual-model architecture, which determines the better model by comparing the prediction results of labeled images with the corresponding Dice loss of the labels. This is more accurate than the previous method of generating pseudo-labels and learning from each other. This paper adopts an architecture where the two models are completely identical. To increase the diversity of results, different network architectures can also be used for the two models.

[0044] (2) Advantages of the Contrastand Complement Block (CCB): This invention proposes a Contrastand Complement Block (CCB), which, while utilizing the higher uncertainty of one model, fully leverages the diversity of prediction results from both models. In regions where predictions differ for the same sample, uncertainty calculation is used to select the prediction result of the other model for complementarity. The prediction results obtained using this method are more stable than those obtained by selecting only a single model or by weighting both models. This module can also be embedded into other semi-supervised learning dual-model approaches.

[0045] (3) Collaborative optimization of the hybrid model architecture: The network consists of two teacher models and one student model. The outputs of the two student models and the teacher model form a diverse supervisory signal to promote model training. This architecture combines the advantages of both student-teacher model architectures and dual-student model architectures. The two model architectures mutually enhance each other, alleviating the problem of parameter coupling between student and teacher models and the model crash caused by abnormal exchange of error information in the dual-student model. It is hoped that they can learn more information by learning their differences, thereby achieving performance improvement.

[0046] (4) Multi-domain experiments verify efficiency: The network was evaluated on datasets from two different domains and compared with several state-of-the-art semi-supervised methods. Experiments show that the proposed scheme achieves the highest efficiency and the network has certain versatility and robustness.

[0047] (5) Wide range of application expansion potential: At the same time, the network model can be used for other 3D medical datasets with good results. Even the dual-student "winner" model architecture and CCB module can be fully embedded into other semi-supervised medical image segmentation networks to bring significant performance gains. Attached Figure Description

[0048] Figure 1 is a schematic diagram of the simulation of the student and teacher model network framework (V-Net) of this invention;

[0049] Figure 2 is a schematic diagram of the structure of the decoupled complementary framework partitioning network (UGCD-Net) of the present invention;

[0050] Figure 3 is a schematic diagram of the complementary module (CCB) of the present invention;

[0051] Figure 4 is a schematic diagram of the process of the semi-supervised medical image decoupling and comparison segmentation method based on uncertainty guidance of the present invention. Detailed Implementation

[0052] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those of ordinary skill in the art without creative effort are within the scope of protection of the present invention.

[0053] The present invention will now be described in further detail with reference to the accompanying drawings and embodiments.

[0054] Example 1, please refer to Figures 1 to 3: This paper proposes an uncertainty-guided decoupling complementarity framework (UGCD) based on semi-supervised medical image segmentation. This network uses two student models and one teacher model for joint training. All three models share the same network architecture, employing a modified V-Net architecture with 5 downsampling layers and 5 upsampling layers. The convolutional kernel size is uniformly 3×3×3, and the activation function is ReLU. See Figure 1 for the specific structure. The modified V-Net achieves good results even with limited data, adapting to 3D medical images while balancing segmentation accuracy and computational efficiency. As shown in Figure 1, the two student models iterate simultaneously, calculating the difference between labeled images and their corresponding labels using the DICE loss. The student model with the better performance is then used to update the teacher model using a moving exponential average. Based on the pseudo-labels provided by the teacher model, the difference between the unlabeled images input to the two student models and the pseudo-labels is calculated to guide parameter updates. To reduce confirmation errors and to reduce erroneous guidance while increasing the diversity of guidance, the two models can learn from each other and exchange their predictions with higher uncertainty. This paper proposes a contrast and complement module, using the output of the better-performing student model calculated using the DICE loss as the primary reference. The dual-student model (CCB) learns from the better-performing model as a reference. In regions where predictions from the two student models at the same location are inconsistent, if the worse-performing model has lower uncertainty, its prediction for that location is adopted. This allows the dual-student model to learn not only the better prediction from both models but also, in regions where predictions are inconsistent but at the same location, the more certain prediction from both models, reducing the exchange of erroneous information and preventing training failures. We evaluated our network on two datasets from different domains and compared it with several state-of-the-art semi-supervised methods. Experiments show that our method achieves the highest performance and that our network has certain generality and robustness.

[0055] As shown in Figure 2, the proposed UGCD-Net architecture consists of two student models and one teacher model, which process labeled and unlabeled images respectively.

[0056] As shown in Figure 2, the lower half contains the labeled image (L in the figure) inputted into the student model. Two student models predict L from the same input. S1 L S2(Representing the outputs of student model 1 and student model 2), the outputs of the two models are then simultaneously input into the comparison module (compare in the figure). In this module, the difference between the model output and the corresponding label is calculated using diceLoss, as shown in the figure below.

[0057]

[0058] In the formula, X refers to the model's output predicted value L. S1 ,L S2 Y refers to the true label corresponding to the model output; the two student models are student model 1 and student model 2.

[0059] In the above formula, X refers to the model's predicted output value L. S1 L S2 Y represents the true label corresponding to the model's output. The model whose predicted result differs most from the true label is called the winner, and the winner updates its parameters to the teacher model using the Moving Exponential Average (EMA) formula. The difference between labeled images and their corresponding true labels is calculated using Dice loss.

[0060] Loss of Student 1 model

[0061] Loss of Student 2 model

[0062] Y represents the actual label value corresponding to the labeled image.

[0063] The upper part of Figure 2 shows the processing of the unlabeled image (U in the figure). First, the prediction result P generated by the teacher model is shown. t , will P t The input is fed into the softmax activation function to obtain the probability values ​​of the output results for the two classes. Then, based on the max function, the index of the class with the higher predicted probability (0 or 1) is determined, which becomes the final pseudo-label for the teacher model (T in the figure). The output U of the two student models is then calculated using the Dice loss function. S1 U S2， The final loss L of the teacher model T T

[0064] Loss of Student Model 1 =

[0065] Loss of Student Model 2 =

[0066] Contrastand Complement Block (CCB)

[0067] First, the uncertainty of the output results of the two models is calculated. The uncertainty H is measured by the entropy of the prediction results, and the calculation formula is as follows:

[0068]

[0069] P represents the probability value of the model's prediction result, i represents the specific application scenario, and i=0 and i=1 represent the two categories of lesion area and non-lesion area.

[0070] HS1 and HS2 represent the uncertainty of the student model's predicted values. One approach is to calculate the uncertainty of the two predictions and then use the prediction category with lower uncertainty from the other model as a pseudo-label to guide the model's learning. The goal is for the two models to learn from each other and acquire more diverse results, making them more robust. However, even when the two models learn from each other and exchange predictions with lower certainty, there's no guarantee that the information exchanged is accurate. The above discussion uses the labeled image and corresponding label DICE loss calculation to determine the "winner" model (the one with the more accurate prediction) and uses the category with lower uncertainty from this model's prediction as a pseudo-label to guide both the model itself and the other model's learning. This approach essentially discards the advantages of the dual-student model, which generates more diverse predictions. Therefore, this paper proposes a Contrastand Complement Block (CCB). First, in regions where both models predict the same outcome and the uncertainty of the predicted outcome is low, the category result is used as a pseudo-label. Second, to increase the diversity of predictions and leverage the predictive advantages of the dual-student model, if the other, less performing model shows lower uncertainty in regions where the two models' predictions are inconsistent, then this prediction result is added to the final dual-model prediction result as the final pseudo-label to guide the two student models in their learning process. The loss from the mutual learning between the two student models is...

[0071] Loss of Student 1 model

[0072] Loss of Student 2 model

[0073] The mask refers to the fact that these two models only calculate predictions with low uncertainty, and the remaining samples are not included in this calculation.

[0074] The total training loss function of UGCD-Net is obtained by using segmentation loss. , L T It is obtained by weighted summation. See below for details:

[0075] Total loss: Loss = + )

[0076] In the formula, Represented as weighting coefficients, Both models were calculated using the formulas above. It is calculated based on the prediction results of labeled images and their corresponding labels. It is calculated using the prediction results of the student model and the pseudo-labels generated by the two student models. It is calculated using the prediction results from the student model and the pseudo-labels generated by the teacher model.

[0077] Example 2; UGCD-Net evaluated on the LA database of the 2018 Atrial Segmentation Challenge. This dataset consists of 100 labeled 3D gadolinium-enhanced left atrial MRI images. 80 of these samples were used for training, and the remaining 20 were used for validation.

[0078] In preprocessing, 3D MR images with enlarged edges are first obtained, and then cropped to the target size as required. They are then further normalized to zero mean and unit variance. During training, the previously obtained images are randomly cropped into 112×112×80 3D patches. Data augmentation is performed using 2D random rotation and flipping operations. The batch size is set to 4, with each batch containing two labeled images and two unlabeled images. The temperature constant T is set to 0.1, and the weights λ are set to a time-varying Gaussian warm-up function. The UGCD-Net is trained for 15K iterations using the SGD optimizer with an initial learning rate of 0.01, decaying by 10% every 2.5K iterations. In testing, a sliding window with a fixed span of 18×18×4 is used to extract features. The predictions from all patches are then recombined into a complete result. All experiments in this paper were conducted on the same equipment environment with a fixed random seed. (Hardware: 14th Gen Intel Core i5-14600kf CPU, NVIDIA GeForce RTX 4070s GPU; Software: PyTorch 1.13.1+cu117, Python 3.9.16).

[0079] Four different metrics were used to evaluate model performance: Dice similarity score (DSC), Jaccard, Hausdorff distance 95 (HD95), and average symmetric distance (ASD). Results for V-Net under various supervision settings (10%, 20%, and fully labeled data) are also provided for reference. Our proposed method achieved state-of-the-art performance on all four evaluation metrics, significantly outperforming the other eight recent state-of-the-art methods, as shown in Table 1. Furthermore, our method achieved a remarkable 91.72% Dice score using only 20% of the labeled data, achieving state-of-the-art segmentation performance on the LA database. For fair comparison, these results were obtained without any post-processing. Experiments show that our proposed method can segment target organs more accurately and precisely, especially in regions prone to misidentification.

[0080] Table 1 shows a comparison with eight state-of-the-art methods on LA data;

[0081]

[0082] The application scenario of this solution is semi-supervised 3D medical image segmentation, especially for the segmentation of 3D MRI images of the left atrium. The semi-supervised medical image segmentation method proposed in this invention has high segmentation accuracy, which is significantly better than the other eight current segmentation methods.

[0083] Please refer to Figures 1 to 3 for the above work process.

[0084] It should be noted that, in this document, relational terms such as "first" and "second" are used merely to distinguish one entity or operation from another, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Furthermore, the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Without further limitations, an element defined by the phrase "comprising one..." does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes said element.

[0085] Although embodiments of the invention have been shown and described, it will be understood by those skilled in the art that various changes, modifications, substitutions and variations can be made to these embodiments without departing from the principles and spirit of the invention, the scope of which is defined by the appended claims and their equivalents.

Claims

1. A semi-supervised medical image decoupling and contrast segmentation method based on uncertainty guidance, characterized in that, The process includes the following steps: S1, constructing a network architecture consisting of two student models and one teacher model, with all three models having the same network architecture; S2, processing the labeled data: inputting labeled medical images into the two student models, constructing a winner student model, which is dynamically determined through Dice loss; the model with the smaller difference between its prediction and the true label is called the winner, and in each iteration, only the winner student model updates the teacher model parameters using the moving exponential average (EMA) to avoid model parameter coupling; the formula for the moving exponential average (EMA) is: In the formula, These are the parameters of the teacher model in the current iteration. These are the teacher model parameters from the previous iteration. S3, Processing unlabeled images: Input the unlabeled medical image into the teacher model to obtain the teacher model's prediction result; Input the prediction result into the softmax activation function to obtain the class probability value, and then determine the class index through the max function, which serves as the pseudo-label for the teacher model; Calculate the difference between the output results of the two student models on the unlabeled medical image and the pseudo-label using the Dice loss function to guide the two student models in updating their parameters; S4, Introduce the Complementary Module (CCB) to calculate the prediction uncertainty of the two student models based on the entropy value, and determine the pseudo-label based on the prediction consistency and uncertainty to guide learning; The formula for calculating the prediction uncertainty of the two student models based on the entropy value in S4 is as follows: In the formula, The entropy value represents the uncertainty of the model's prediction; P represents the probability of the model's prediction; and i represents the number of categories. First, in regions where the two models predict the same result and the uncertainty of the prediction is low, the category result is used as a pseudo-label. Second, in order to increase the diversity of predictions and combine the prediction advantages of the dual-student model, if the other, less effective model shows lower uncertainty in regions where the two models' predictions are inconsistent, then this prediction result will be added to the final dual-model prediction result as the final pseudo-label to guide the two student models in their learning. S5, calculate the loss according to the total loss formula and iteratively train the network until convergence.

2. The semi-supervised medical image decoupling and comparison segmentation method based on uncertainty guidance according to claim 1, characterized in that: The formula for calculating the Dice loss function described in S2 is as follows: In the formula, X refers to the model's output predicted value L. S1 L S2 Y refers to the true label corresponding to the model output; the two student models are student model 1 and student model 2, and the Dice loss formula for student model 1 is: In the formula, L S1 Let Y represent the predicted output probability value of student model 1; let Y represent the true label value corresponding to the labeled image; the Dice loss formula for student model 2 is: In the formula, L S2 Y represents the predicted output probability value of Student Model 2; Y represents the true label value corresponding to the labeled image.

3. The uncertainty-guided semi-supervised medical image decoupling and comparison segmentation method according to claim 2, characterized in that: The formula for calculating the loss of student model 1 and pseudo-labels in S3 is: = In the formula, Let T represent the output of student model 1 on unlabeled medical images; T represents the pseudo-labels generated by the teacher model; the formula for calculating the loss of student model 2 and the pseudo-labels is: = In the formula, T represents the output of student model 2 on unlabeled medical images; T represents the pseudo-labels generated by the teacher model.

4. The uncertainty-guided semi-supervised medical image decoupling and comparison segmentation method according to claim 3, characterized in that: In S4, the formula for calculating the mutual learning loss of student model 1 is: The formula for calculating the mutual learning loss of student model 2 is: In the formula, "mask" refers to the fact that the two models only calculate the prediction results with low uncertainty, and the remaining samples are not included in this calculation.

5. The semi-supervised medical image decoupling and comparison segmentation method based on uncertainty guidance according to claim 1, characterized in that: In S5, the formula for calculating the total loss is: Loss = + In the formula, Represented as weighting coefficients, It is calculated based on the prediction results of labeled images and their corresponding labels. It is calculated using the prediction results of the student model and the pseudo-labels generated by the two student models. It is calculated using the prediction results from the student model and the pseudo-labels generated by the teacher model.

6. The semi-supervised medical image decoupling and comparison segmentation method based on uncertainty guidance according to claim 1, characterized in that: It is used in medicine to accurately segment abnormal regions in magnetic resonance images, assisting doctors in segmenting lesions.

Images