A Medical Image Segmentation Method Based on Uncertainty Estimation and Multi-level Distillation

Abstract

The present invention discloses a medical image segmentation method based on uncertainty estimation and multi-stage distillation, which relates to medical image segmentation. The method comprises quantifying the prediction uncertainty of a teacher model for an input medical image through an uncertainty estimation module to generate uncertainty weights; screening high-confidence pseudo-labels based on the uncertainty weights, and transferring the knowledge of the teacher model to a student model in stages through a multi-stage knowledge distillation module; embedding a hybrid attention module in the teacher and student models, combining a channel attention mechanism with a spatial attention mechanism to enhance feature capture of lesion structure and boundaries; optimizing the student model training process using a composite loss function that fuses cross entropy loss, Dice loss, and uncertainty-based distillation loss; and applying the trained network to a medical image segmentation scenario to complete medical image segmentation based on uncertainty estimation and multi-stage distillation. The present invention solves the problem of scarcity of high-quality annotated data in medical image segmentation.

Description

Technical Field

[0001] The present invention relates to medical image segmentation, and in particular to a medical image segmentation method based on uncertainty estimation and multi-stage distillation. Background Art

[0002] Medical image segmentation has been widely used in clinical practice, assisting doctors in providing high-quality diagnoses by identifying lesions or organs. Deep learning has seen significant development in recent years, significantly improving the performance of medical image segmentation. Currently, mainstream deep learning segmentation networks include convolutional neural networks, the attention-based Transformer model, and hybrid models that combine the strengths of various architectures. CNNs extract local features from medical images through hierarchical convolution operations, exhibiting strong spatial invariance and feature sharing capabilities, making them excellent for medical image segmentation tasks and particularly suitable for structural analysis sensitive to local details. In contrast, the Transformer model leverages a self-attention mechanism to model long-range dependencies, effectively capturing global information. This allows for stronger modeling capabilities in segmenting complex morphological structures and improves the accuracy of boundary regions and small objects. However, regardless of the network model used, performance improvements typically rely on large amounts of labeled data, which is a significant challenge for medical image segmentation. The annotation process for medical images requires not only highly specialized medical knowledge but also a significant amount of time to ensure the accuracy of segmentation labels, making the acquisition of high-quality labeled data extremely expensive. To solve this problem, many studies in recent years have been devoted to exploring how to improve segmentation performance under the condition of limited labeled data.

[0003] Semi-supervised learning methods can improve the learning performance of the model by effectively utilizing unlabeled data when there is limited labeled data and a large amount of unlabeled data. Specifically, semi-supervised learning combines a small number of labeled samples with a large number of unlabeled samples, and enables the model to extract valuable features from unlabeled data by generating pseudo labels, contrastive learning, or consistency training. Currently, semi-supervised learning (SSL) methods can be roughly divided into two categories. The first category is pseudo-label-based methods, which generate pseudo labels through the model's prediction of unlabeled data, and use these pseudo labels together with labeled data for training. The other category is consistency regularization-based methods, which use two deep convolutional neural networks to learn the consistency of inference results of the same unlabeled image under different perturbation conditions, thereby better utilizing the information of unlabeled data.

[0004] However, uncertainty is prevalent in medical datasets used to train deep learning algorithms due to interference from imaging equipment, individual patient differences, and errors in manual annotation. This uncertainty severely limits the performance of image segmentation algorithms. Summary of the Invention

[0005] In response to the above-mentioned deficiencies in the prior art, the present invention provides a medical image segmentation method based on uncertainty estimation and multi-stage distillation to solve the problem of scarcity of high-quality labeled data in medical image segmentation.

[0006] In order to achieve the above-mentioned object of the invention, the technical solution adopted by the present invention is: a medical image segmentation method based on uncertainty estimation and multi-stage distillation, comprising the following steps:

[0007] S1: Quantify the prediction uncertainty of the teacher model for the input medical image through the uncertainty estimation module and generate uncertainty weights;

[0008] S2: Filter high-confidence pseudo-labels based on uncertainty weights and transfer the knowledge of the teacher model to the student model in stages through a multi-stage knowledge distillation module;

[0009] S3: Embed a hybrid attention module in the teacher model and the student model, combining the channel attention mechanism and the spatial attention mechanism to enhance the feature capture of the lesion structure and boundary;

[0010] S4: A composite loss function that integrates cross entropy loss, Dice loss, and uncertainty-based distillation loss is used to optimize the training process of the student model.

[0011] S5: Use the trained network in medical image segmentation scenarios to complete medical image segmentation based on uncertainty estimation and multi-level distillation.

[0012] Furthermore, the uncertainty estimation module in S1 is implemented by the Monte Carlo Dropout method, which specifically includes:

[0013] A1: Enable the Dropout layer during the teacher model inference phase and make multiple forward predictions for the same input;

[0014] A2: Calculate the prediction entropy based on the probability distribution of multiple prediction results as a quantitative indicator of segmentation uncertainty:

[0015]

[0016]

[0017] in, To predict entropy, is the number of classes in the segmentation task, is the average softmax probability of T samples from the teacher model, is the teacher model in the tth forward pass The predicted value of the class;

[0018] A3: Dynamically adjust the pseudo-label screening threshold based on the predicted entropy to filter the prediction results in high uncertainty areas.

[0019] Furthermore, the implementation of the multi-level knowledge distillation module in S2 includes:

[0020] B1: Update the weights between the teacher model and the student model through exponential moving average to generate an intermediate model;

[0021] B2: Use the intermediate model as a new teacher model for secondary distillation, transferring multi-level semantic features step by step;

[0022] B3: During the distillation process, the uncertainty map based on the prediction entropy is used to optimize the mean squared error loss and Dice loss between the student model's prediction and the teacher model's prediction. At the same time, the cross entropy loss and Dice loss between the student model and the binary label are calculated to constrain the prediction consistency of the student model and the teacher model.

[0023] Furthermore, the network architectures of the teacher model and the student model in S3 are both 3D Unet, and a hybrid attention module is embedded in the bridge part of the 3D Unet. The hybrid attention module includes the following operations:

[0024] C1: Perform global average pooling and global maximum pooling on the input feature map to generate two channel descriptors:

[0025]

[0026]

[0027] in, and For two channel descriptors, is the number of channels, and are the height and width of the feature map, respectively. and Represents the first Row and Elements of the column;

[0028] C2: Generate channel attention weights by interacting the two channel descriptors through a multi-layer perceptron:

[0029]

[0030] in, is the channel attention weight, is a multilayer perceptron;

[0031] C3: Multiply the channel attention weight by the input feature map element by element to obtain the channel-enhanced feature map, perform spatial dimension pooling on the channel-enhanced feature map, and generate spatial attention weights through convolution operation. The channel-enhanced feature map is expressed as:

[0032]

[0033] in, It is the feature map after channel enhancement.

[0034] Furthermore, the composite loss function in S4 for:

[0035]

[0036] in, is the cross entropy loss, is the Dice loss, is the uncertainty distillation loss, and is the balance coefficient.

[0037] Furthermore, the cross entropy loss is:

[0038]

[0039] in, is the total number of samples, For the student model to the current sample predictions, For the current sample The actual label.

[0040] Furthermore, the Dice loss is:

[0041]

[0042] in, is the intersection symbol.

[0043] Furthermore, the distillation loss of the uncertainty is:

[0044]

[0045] in, is the time-dependent Gaussian warm-up function, is the number of iterations currently reached during training, The teacher model for the current sample The segmentation uncertainty, is the threshold of the sample that the teacher model is most certain of, is a hyperparameter, For the student model to the current sample Prediction And the teacher model for the current sample Prediction The mean square error loss between For the student model to the current sample Prediction And the teacher model for the current sample Prediction The Dice loss between Indicates selection operation, if If true, the value is 1, otherwise the value is 0;

[0046]

[0047]

[0048] in, is a single voxel, is the total number of voxels in the current sample, For the student model, Sample No. The predicted value of a voxel, The teacher model is the current Sample No. The predicted value of a voxel.

[0049] The beneficial effects of the present invention are:

[0050] To address the scarcity of high-quality annotated data in medical image segmentation, the present invention proposes a segmentation network (UGMLDS-Net) based on uncertainty estimation and multi-level knowledge distillation. First, by introducing an uncertainty estimation module, the method utilizes the Monte Carlo Dropout method to quantify the prediction uncertainty of the teacher model, thereby dynamically screening high-confidence pseudo-labels and reducing the negative impact of noisy pseudo-labels on model training. In addition, UGMLDS-Net adopts a multi-level knowledge distillation mechanism, which enables the student model to gradually learn the multi-level and multi-stage information in the teacher model through staged knowledge transfer, thereby more accurately capturing the details and boundaries of the lesions. At the same time, in order to enhance the model's ability to capture complex lesion structures and fuzzy boundaries, the present invention inserts a hybrid attention module into the 3DUnet network. This module combines channel attention and spatial attention mechanisms to significantly improve segmentation accuracy.

[0051] The UGMLDS-Net proposed in this paper combines uncertainty estimation and multi-level knowledge distillation mechanism to achieve finer-grained knowledge transfer, enhance the learning ability of the student model, and realize efficient utilization of unlabeled data under the condition of limited labeled data, providing an effective solution to the problem of scarce labeled data in medical image segmentation.

[0052] This paper proposes an uncertainty estimation module to guide the student model to perform more refined learning in high uncertainty areas, optimize the distillation process, and enhance the model's ability to distinguish difficult-to-segment areas.

[0053] This paper proposes an uncertainty-based loss function, which dynamically filters the knowledge of the teacher model during the multi-level distillation process, adjusts the knowledge transfer through uncertainty weights, enables the student model to preferentially learn information in high-confidence areas, and performs progressive knowledge optimization under the multi-level distillation framework, thereby improving network performance. BRIEF DESCRIPTION OF THE DRAWINGS

[0054] Figure 1 This is a flow chart of a medical image segmentation method based on uncertainty estimation and multi-stage distillation of the present invention.

[0055] Figure 2 This is a structural diagram of a medical image segmentation network based on uncertainty estimation and multi-stage distillation in the present invention.

[0056] Figure 3 A visual comparison of segmentation using 10% labeled data on the Brats 19 dataset.

[0057] Figure 4 A visual comparison of segmentation using 20% labeled data on the LA dataset. DETAILED DESCRIPTION

[0058] The present invention will be further described below with reference to the accompanying drawings and specific embodiments.

[0059] like Figure 1 As shown, a medical image segmentation method based on uncertainty estimation and multi-level distillation includes the following steps:

[0060] S1: Quantify the prediction uncertainty of the teacher model for the input medical image through the uncertainty estimation module and generate uncertainty weights;

[0061] S2: Filter high-confidence pseudo-labels based on uncertainty weights and transfer the knowledge of the teacher model to the student model in stages through a multi-stage knowledge distillation module;

[0062] S3: Embed a hybrid attention module in the teacher model and the student model, combining the channel attention mechanism and the spatial attention mechanism to enhance the feature capture of the lesion structure and boundary;

[0063] S4: A composite loss function that integrates cross entropy loss, Dice loss, and uncertainty-based distillation loss is used to optimize the training process of the student model.

[0064] S5: Use the trained network in medical image segmentation scenarios to complete medical image segmentation based on uncertainty estimation and multi-level distillation.

[0065] Currently, most deep learning-based medical image segmentation algorithms are data-driven, and their performance is highly dependent on the quality of the dataset. Tumors often exhibit significant morphological and histological heterogeneity. Different internal structures (such as necrotic areas, edema, and enhancement cores), as well as tumor shape, size, and location, can vary significantly between patients. Furthermore, due to differences in resolution, contrast, and noise levels among imaging devices, acquired MRI images may exhibit issues such as modality loss or image artifacts. Furthermore, when constructing datasets for deep learning model training, medical images must be annotated by professional physicians. However, due to the subjectivity of expert judgment of tumor boundaries, annotation results can be inconsistent due to inter-expert variability. In summary, the uncertainty introduced into tumor segmentation datasets due to individual patient differences, tumor imaging noise, and expert annotation variability poses a significant challenge to accurate medical image segmentation. Therefore, this paper proposes an uncertainty quantification method that quantifies the uncertainty in MRI images using a teacher model. This uncertainty information guides a student model to achieve more refined segmentation of tumor regions.

[0066] Bayesian methods quantify uncertainty by specifying a prior distribution of neural network parameters and computing the posterior distribution of the parameters given training data. Given a dataset and , the random output of the Bayesian network is expressed as , the model likelihood is expressed as , Bayesian inference is used to calculate the posterior For the weight w, given the prior Capture a reasonable set of model parameters:

[0067]

[0068] Given a new test sample , the predictive distribution can be obtained by integrating the model parameters according to their posterior distribution:

[0069]

[0070] For MRI images, the above formula can be used to generate a posterior weighted average of the model for each pixel in the image, also known as the Bayesian model average (BMA). However, the difficulty with the above predictive distribution is that the posterior of the parameters is difficult to calculate analytically:

[0071]

[0072] The above formula needs to be integrated over the space of all possible weights in order to be marginalized, but deep neural networks typically have millions of parameters, making it impossible to integrate over all weighted parameters. This leads to approximate Bayesian methods for approximating the posterior.

[0073] Since the MCDropout method only requires adding a dropout layer to the existing network model without changing its architecture, it has become a common method for uncertainty quantification. Therefore, this embodiment uses the MCDropout method to quantify the uncertainty information in remote sensing images. During the training process, the dropout layer randomly discards some neural network units to avoid overfitting.

[0074] Most existing image segmentation algorithms give a certain probability score. Using the MCDropout method, when training the model, we enable the dropout layer to perform multiple forward predictions on the same remote sensing image during testing, which is equivalent to obtaining samples from the posterior. In the experiment, we do this 50 times to obtain 50 Monte Carlo samples that approximate the predicted probability distribution.

[0075] Based on samples drawn using Monte Carlo methods, the prediction entropy can be quantified to approximate the segmentation uncertainty of the teacher model.

[0076] The uncertainty estimation module in S1 is implemented by the Monte Carlo Dropout method, which specifically includes:

[0077] A1: Enable the Dropout layer during the teacher model inference phase and make multiple forward predictions for the same input;

[0078] A2: Calculate the prediction entropy based on the probability distribution of multiple prediction results as a quantitative indicator of segmentation uncertainty:

[0079]

[0080]

[0081] in, To predict entropy, is the number of classes in the segmentation task, is the average softmax probability of T samples from the teacher model, is the teacher model in the tth forward pass The predicted value of the class;

[0082] A3: Dynamically adjust the pseudo-label screening threshold based on the predicted entropy to filter the prediction results in high uncertainty areas.

[0083] For a given input, the samples obtained by multiple forward predictions of the teacher model can estimate the overall segmentation uncertainty Under the guidance of the uncertainty of the teacher model, the teacher model filters out predictions with high uncertainty that are considered relatively unreliable, and only retains the teacher model's relatively certain predictions for the student model to learn. This can, to a certain extent, avoid the transfer of erroneous information from the teacher model during the student model's refined segmentation of the lesion area, thereby improving the segmentation accuracy of the tumor area.

[0084] Knowledge distillation was first proposed as a network compression method. By allowing a lightweight student model to learn the output of the teacher model, the student network can achieve better results than learning directly on its own. After continuous research and development, a variety of methods for learning teacher model knowledge have emerged, such as multi-teacher distillation, self-distillation, and online distillation. Although the existing knowledge distillation mechanism has achieved good results in medical image segmentation, it still has certain shortcomings in semi-supervised scenarios. Existing knowledge distillation methods usually rely only on one-time knowledge transfer, that is, the knowledge transfer of the teacher model to the student model is completed through a single distillation process. However, in complex medical image segmentation, this method is difficult to fully utilize the multi-level and multi-stage information contained in the teacher model. It is difficult for the student model to fully grasp the multi-scale structure and fine-grained structural information of the lesion after a single distillation. Therefore, the present invention proposes a multi-level knowledge distillation mechanism based on the above-mentioned uncertainty quantification module, such as Figure 2 As shown in Figure 1. First, the yellow teacher model at the top and the blue intermediate model in the middle are trained. This is the first knowledge distillation. After training, the intermediate model is used as the teacher model and trained together with the green student model at the bottom. This is the second knowledge distillation. All three models use 3D UNet as the underlying architecture, embedded with the hybrid attention module described in C1. The entire processing logic can be divided into two parts. The first distillation trains the top and middle layers of the network simultaneously, while the second distillation trains the middle and bottom layers simultaneously. Figure 2 The output black and white predictions of the intermediate model in the middle blue are , and it serves as the output of the student model in the first distillation, and it serves as the teacher model in the second distillation, so its output at this time is .

[0085] In multi-level knowledge distillation, the teacher model (TeacherModel) first undergoes knowledge distillation to obtain an intermediate model (MediateModel), and then the intermediate model (MedicateModel) is used as the teacher model for distillation again to obtain the final student model (StudentModel). Specifically, the teacher model (TeacherModel), the intermediate model (MediateModel), and the student model (StudentModel) all use the same network architecture. During each knowledge distillation training process, we will Update to student model weights The exponential moving average of . Where γ is a hyperparameter used for weight decay. During the knowledge distillation process, the uncertainty graph based on the predicted entropy is used to optimize the mean square error loss and Dice loss between the predictions of the student model and the teacher model. At the same time, the cross entropy loss and Dice loss between the student model and the binary label are calculated, so that the student model can not only learn the information in the label, but also learn the more certain knowledge of the teacher model through the guidance of uncertainty. In addition, through multi-level knowledge distillation, the multi-level and multi-stage information contained in the teacher model can be fully utilized to better grasp the fine-grained information of the lesion and refine the segmentation accuracy of the lesion boundary.

[0086] The implementation of the multi-level knowledge distillation module in S2 includes:

[0087] B1: Update the weights between the teacher model and the student model through exponential moving average to generate an intermediate model;

[0088] B2: Use the intermediate model as a new teacher model for secondary distillation, transferring multi-level semantic features step by step;

[0089] B3: During the distillation process, the uncertainty map based on the prediction entropy is used to optimize the mean squared error loss and Dice loss between the student model's prediction and the teacher model's prediction. At the same time, the cross entropy loss and Dice loss between the student model and the binary label are calculated to constrain the prediction consistency of the student model and the teacher model.

[0090] The network architectures of the teacher model and the student model in S3 are both 3D Unet, and a hybrid attention module is embedded in the bridge part of the 3D Unet. The hybrid attention module includes the following operations:

[0091] C1: Perform global average pooling and global maximum pooling on the input feature map to generate two channel descriptors:

[0092]

[0093]

[0094] in, and For two channel descriptors, is the number of channels, and are the height and width of the feature map, respectively. and Represents the first Row (height) and Column (width) elements;

[0095] C2: Generate channel attention weights by interacting the two channel descriptors through a multi-layer perceptron:

[0096]

[0097] in, is the channel attention weight, is a multilayer perceptron;

[0098] C3: Multiply the channel attention weight by the input feature map element by element to obtain the channel-enhanced feature map, perform spatial dimension pooling on the channel-enhanced feature map, and generate spatial attention weights through convolution operation. The channel-enhanced feature map is expressed as:

[0099]

[0100] in, It is the feature map after channel enhancement.

[0101] In this embodiment, both the teacher and student networks use a 3D Unet architecture. A hybrid attention module (CBAM) is inserted into the central bridge portion of the 3D Unet to enhance its expressive power. The CBAM module combines the channel attention mechanism and the spatial attention mechanism to better capture key image features and improve model performance. The spatial attention mechanism focuses on the spatial dimension of the image, generating an importance map for each spatial region by performing max pooling and average pooling on the input feature map. The pooling results are then fused through convolution operations to generate spatial attention weights for each location. This allows the model to focus on key regions in the spatial dimension, particularly boundary regions. The insertion of the hybrid attention module facilitates multi-scale feature extraction. The channel attention mechanism selectively enhances important channels, while the spatial attention mechanism focuses on key regions, providing a more detailed feature representation. CBAM can be easily integrated into existing network architectures without adding excessive computational overhead, while significantly improving model performance. For complex structures and subtle differences in images, CBAM helps the model effectively capture key information, achieving particularly good results in segmenting boundaries and details.

[0102] The composite loss function in S4 for:

[0103]

[0104] in, is the cross entropy loss, is the Dice loss, is the uncertainty distillation loss, and is the balance coefficient.

[0105] During the knowledge distillation training process, by calculating the cross entropy loss and Dice loss between the student model's prediction and the label, as well as the mean square error loss and Dice loss between the student model's prediction and the teacher model's prediction, the student model can not only learn the lesion information in the label, but also learn the relatively correct and reliable knowledge in the teacher model. Through multiple knowledge distillations, the student model's segmentation accuracy for the lesion area and lesion boundary can be effectively improved. In this embodiment and Take 0.5 and 1 respectively.

[0106] Cross entropy loss calculation for student model predictions With label The cross entropy loss.

[0107] The cross entropy loss is:

[0108]

[0109] in, is the total number of samples, For the student model to the current sample predictions, For the current sample The actual label.

[0110] In medical image segmentation, cross-entropy loss primarily optimizes the model by comparing predictions and labels pixel by pixel, making it well-suited for scenarios with balanced class distributions. However, medical image segmentation often faces a serious class imbalance problem, where lesions or organ regions only account for a small proportion of the image. This inter-class imbalance causes cross-entropy loss to primarily optimize background regions, thus affecting the accurate segmentation of foreground regions. Therefore, the present invention uses the Dice loss function to assist in model optimization.

[0111] The Dice loss function evaluates model performance by calculating the Dice coefficient between the student model's predictions and the true labels. The Dice loss function avoids the dominance of excessive background pixels in loss calculations and is therefore suitable for segmenting small lesions. By focusing on the global optimization goal of the overall degree of match between model predictions and true labels, rather than independent pixel-by-pixel calculations, the Dice loss function can improve the overall consistency of segmentation results and make model predictions more accurate and smoother. By introducing the Dice loss function on top of the cross-entropy loss function, the model's lesion segmentation performance can be effectively improved in scenarios where the lesion ratio is relatively small, thereby improving the overall quality of the lesion segmentation results.

[0112] The Dice loss is:

[0113]

[0114] in, is the intersection symbol.

[0115] In addition, the present invention also introduces uncertainty loss to filter the knowledge of the teacher model during the multi-level knowledge distillation process, guiding the student model to gradually learn the correct knowledge in the teacher model over multiple learning cycles. The proposed uncertainty loss is as follows:

[0116] The distillation loss of the uncertainty is:

[0117]

[0118] in, is the time-dependent Gaussian warm-up function, is the number of iterations currently reached during training, The teacher model for the current sample The segmentation uncertainty, is the threshold of the sample that the teacher model is most certain of, is a hyperparameter, For the student model to the current sample Prediction And the teacher model for the current sample Prediction The mean square error loss between For the student model to the current sample Prediction And the teacher model for the current sample Prediction The Dice loss between Indicates selection operation, if If true, the value is 1, otherwise the value is 0;

[0119] In the above uncertainty loss function, , Indicates the maximum number of iterations set. In this embodiment, the same Gaussian ramp-up function is used during the experiment. By the formula H=(0.75+0.25×λ(t))× The threshold H is changed from rise to . Represents the maximum uncertainty value, which is set to ln(2) in the experiment. In the experiment, it is set to 0.75;

[0120]

[0121]

[0122] in, is a single voxel, is the total number of voxels in the current sample, For the student model, Sample No. The predicted value of a voxel, The teacher model is the current Sample No. The predicted value of a voxel.

[0123] The mean square error loss can capture the fine-grained differences between the predictions of the teacher model and the student model, and make the student model closer to the output of the teacher model. In the case of a serious imbalance between the lesion area and the background area, the Dice loss helps the student model better learn and optimize the teacher model's processing and segmentation of the lesion area boundary. In the process of training the student model, the teacher model's segmentation threshold H for the sample is changed from rise to , so that the student model learns from relatively certain examples to relatively unreliable samples. In addition, by using By adjusting the learning weights of mean square error loss and Dice loss through hyperparameters, the student model can stabilize the numerical regression while optimizing the teacher model’s prediction of the shape and boundary of the lesion area, thereby achieving an overall improvement in segmentation accuracy.

[0124] In one embodiment of the present invention, the experimental datasets in this paper adopt the BraTS2019

[47] and LA

[48] datasets. BraTS is an important public dataset for multimodal brain tumor segmentation and is widely used in the research of this topic. The BraTS2019 dataset contains MRI samples of 335 glioma patients with expert annotations. Each patient's MRI scan contains four different modalities, including T1, T1ce, T2 and Flair. We use the FLAIR modality for tumor segmentation because this modality can better show malignant tumors. In the experiment, this paper resampled all MRI scans to the same resolution size and normalized the intensity to zero mean and unit variance. In the experiment, the BraTS2019 dataset was divided into 250 scans for training, 25 scans as validation set, and the remaining 60 scans as test set. Among the 250 training scans, 10% of the 25 scans and 20% of the 50 scans were used as labeled data according to different settings, and the remaining scans were used as unlabeled data. The LA dataset is a public dataset focused on left atrial segmentation, containing numerous expert-labeled left atrial images. In this study, the dataset was divided into a training set of 80 scans and a validation set of 20 scans. The network was trained using 16 scans, or 20% of the labeled data.

[0125] In order to ensure a fair comparison of the experimental results, this paper uses V-Net

[26] as the backbone structure. In order to control the balance between the supervised segmentation loss and the unsupervised consistency loss, a Gaussian ramp function is used in the experiment. . The stochastic gradient descent (SGD) optimizer is used to update the network parameters, with an initial learning rate of 1e-2 and a decay of 0.1 every 2500 iterations. The maximum number of training iterations is set to 6000. The batch size is set to 4, and each mini-batch contains 2 labeled images and 2 unlabeled images. This experiment randomly crops 112×112×80 sub-volumes as network input, and uses a sliding window strategy to obtain the final segmentation results. The experiment uses standard data augmentation techniques during training to avoid overfitting, including random flipping, and rotations of 90 degrees, 180 degrees, and 270 degrees along the axial plane. In order to quantitatively evaluate the segmentation results, four complementary evaluation metrics are used. Two metrics, the region-based Dice similarity coefficient (Dice) and the Jaccard index (Jaccard), are used to measure region mismatch. The average surface distance (ASD) and the 95% Hausdorff distance (95HD), two boundary-based metrics, are used to evaluate the boundary error between the segmentation results and the ground truth.

[0126] Tables 1 and 2 show the quantitative comparison results of the proposed method with the current mainstream semi-supervised medical image segmentation methods on the BraTS2019 dataset, including baseline models such as MT (MeanTeachers), UA-MT (uncertainty-aware self-ensembling mean teacher framework), URPC (Uncertainty Rectified Pyramid Consistency), DTC (Dual-task Consistency), and UG-MCL (Uncertainty-guided mutual consistency learning).

[0127] Table 1 Comparison of the performance of each method on the BraTS2019 dataset using 10% labeled data

[0128]

[0129] Table 2 Comparison of the performance of each method on the BraTS 2019 dataset using 20% labeled data

[0130]

[0131] The results in Table 1 show that the proposed UGMLDS-Net (Uncertainty-Guided Multi-Level Distillation Segmentation Network) achieves optimal results across all four metrics. Specifically, in the core metric of Dice similarity coefficient, UGMLDS-Net achieves 85.21%, approximately 2% higher than the second-ranked method, UG-MCL. This improvement demonstrates that UGMLDS-Net has a significant advantage in overall lesion segmentation accuracy, enabling more accurate identification and segmentation of complex lesion structures. In terms of 95HD, a key metric measuring segmentation boundary error, UGMLDS-Net performs even better, achieving a result of 8.41, a significant 35% improvement compared to UG-MCL's 11.44 pixels. This demonstrates that UGMLDS-Net not only excels in overall segmentation accuracy but also demonstrates a significant advantage in boundary capture, enabling more precise localization of lesion edges. UGMLDS-Net also demonstrates comprehensive performance advantages in the other two metrics. These results show that UGMLDS-Net outperforms other methods in segmentation accuracy, boundary capturing ability, and robustness.

[0132] UGMLDS-Net also demonstrated comprehensive performance advantages in the other two metrics. These results demonstrate that UGMLDS-Net outperforms other methods in segmentation accuracy, boundary capture, and robustness. Table 3 further shows the performance of UGMLDS-Net on the left atrial segmentation dataset. Experimental results demonstrate that even with only 20% labeled data, UGMLDS-Net can still achieve optimal segmentation results. This further validates the versatility and effectiveness of UGMLDS-Net in various medical image segmentation tasks.

[0133] Table 3 Performance comparison of each method on the LA dataset using 20% labeled data

[0134]

[0135] In order to more intuitively demonstrate the performance advantages of UGMLDS-Net, Figure 3 and 4Visualization results are presented for two representative brain regions and the left atrium segmentation task. Comparing the segmentation performance of different methods in these key regions clearly demonstrates the superior ability of UGMLDS-Net in handling complex lesion structures and fuzzy boundary regions. Specifically, in the core of the lesion, UGMLDS-Net demonstrates superior detail capture, accurately identifying and segmenting the lesion's subtle structure. In the edge regions, its boundary segmentation is smoother and more accurate, effectively avoiding over-segmentation or under-segmentation. This advantage is consistent with the significant improvement in the Dice coefficient and 95HD metric in the experiments, further validating the superiority of UGMLDS-Net in semi-supervised medical image segmentation tasks. These results demonstrate that UGMLDS-Net not only performs well in quantitative metrics but also has significant advantages in actual segmentation results, providing more reliable and accurate segmentation results for clinical diagnosis. In summary, by combining uncertainty estimation and a multi-level knowledge distillation mechanism, UGMLDS-Net achieves efficient utilization of both labeled and unlabeled data in semi-supervised medical image segmentation tasks, achieving state-of-the-art performance across all evaluation metrics. This method provides an effective solution to the problem of scarce labeled data in medical image segmentation.

[0136] Those skilled in the art will appreciate that the embodiments described herein are intended to help readers understand the principles of the present invention, and it should be understood that the scope of protection of the present invention is not limited to such specific descriptions and embodiments. Those skilled in the art can make various other specific variations and combinations based on the technical teachings disclosed in the present invention without departing from the essence of the present invention, and such variations and combinations are still within the scope of protection of the invention.

Claims

1. A medical image segmentation method based on uncertainty estimation and multi-stage distillation, characterized in that: The following steps are involved: S1: Quantify the prediction uncertainty of the teacher model for the input medical image through the uncertainty estimation module and generate uncertainty weights; S2: Filter high-confidence pseudo-labels based on uncertainty weights and transfer the knowledge of the teacher model to the student model in stages through a multi-stage knowledge distillation module; The implementation of the multi-level knowledge distillation module in S2 includes: B1: Update the weights between the teacher model and the student model through exponential moving average to generate an intermediate model; B2: Use the intermediate model as a new teacher model for secondary distillation, transferring multi-level semantic features step by step; B3: During the distillation process, the uncertainty graph based on the prediction entropy is used to optimize the mean squared error loss and Dice loss between the student model's prediction and the teacher model's prediction. At the same time, the cross entropy loss and Dice loss between the student model and the binary label are calculated to constrain the prediction consistency of the student model and the teacher model. S3: Embed a hybrid attention module in the teacher model and the student model, combining the channel attention mechanism and the spatial attention mechanism to enhance the feature capture of the lesion structure and boundary; S4: A composite loss function that integrates cross entropy loss, Dice loss, and uncertainty-based distillation loss is used to optimize the training process of the student model. The composite loss function in S4 for: ; in, is the cross entropy loss, is the Dice loss, is the uncertainty distillation loss, and is the balance coefficient; The distillation loss of the uncertainty is: ; in, is the time-dependent Gaussian warm-up function, is the number of iterations currently reached during training, The teacher model for the current sample The segmentation uncertainty, is the threshold of the sample that the teacher model is most certain of, is a hyperparameter, For the student model to the current sample Prediction And the teacher model for the current sample Prediction The mean square error loss between For the student model to the current sample Prediction And the teacher model for the current sample Prediction The Dice loss between Indicates selection operation, if If true, the value is 1, otherwise the value is 0. is the total number of samples; ; ; in, is a single voxel, is the total number of voxels in the current sample, For the student model, Sample No. The predicted value of a voxel, The teacher model is the current Sample No. The predicted value of the voxel; S5: Use the trained network in medical image segmentation scenarios to complete medical image segmentation based on uncertainty estimation and multi-level distillation.

2. The medical image segmentation method based on uncertainty estimation and multi-stage distillation according to claim 1, characterized in that: The uncertainty estimation module in S1 is implemented by the Monte Carlo Dropout method, which specifically includes: A1: Enable the Dropout layer during the teacher model inference phase and make multiple forward predictions for the same input; A2: Calculate the prediction entropy based on the probability distribution of multiple prediction results as a quantitative indicator of segmentation uncertainty: ; ; in, To predict entropy, is the number of classes in the segmentation task, is the average softmax probability of T samples from the teacher model, is the teacher model in the tth forward pass The predicted value of the class; A3: Dynamically adjust the pseudo-label screening threshold based on the predicted entropy to filter the prediction results in high uncertainty areas.

3. The medical image segmentation method based on uncertainty estimation and multi-stage distillation according to claim 1, characterized in that: The network architectures of the teacher model and the student model in S3 are both 3D Unet, and a hybrid attention module is embedded in the bridge part of the 3D Unet. The hybrid attention module includes the following operations: C1: Perform global average pooling and global maximum pooling on the input feature map to generate two channel descriptors: ; ; in, and For two channel descriptors, is the number of channels, and are the height and width of the feature map, respectively. and Represents the first Row and Elements of the column, is the input feature map; C2: Generate channel attention weights by interacting the two channel descriptors through a multi-layer perceptron: ; in, is the channel attention weight, is a multilayer perceptron; C3: Multiply the channel attention weight by the input feature map element by element to obtain the channel-enhanced feature map, perform spatial dimension pooling on the channel-enhanced feature map, and generate spatial attention weights through convolution operations. The channel-enhanced feature map is expressed as: ; in, It is the feature map after channel enhancement.

4. The medical image segmentation method based on uncertainty estimation and multi-stage distillation according to claim 1, characterized in that: The cross entropy loss is: ; in, is the total number of samples, For the student model to the current sample predictions, For the current sample The actual label.

5. The medical image segmentation method based on uncertainty estimation and multi-stage distillation according to claim 4, characterized in that: The Dice loss is: ; in, is the intersection symbol.

Images