Uncertainty-aware collaborative learning with mixed images for semi-supervised medical image segmentation

Semi-supervised medical image segmentation (SSMIS) addresses the high acquisition cost of medical annotation by using unlabeled data. Pseudo-labeling (PL) is a mainstream method for tackling unlabeled data in SSMIS. A key challenge for PL-based methods is models trained with limited annotation data would generate noisy pseudo-labels, which provide unreliable supervision to degrade the learning quality. Additionally, perturbations are introduced to improve the performance of PL-based methods. Most methods use a single type of perturbation, which limits the ability of the models to handle more unknown data. Employing multiple types of perturbations can address the above limitation, but ensuring training quality remains a challenge under diverse perturbations. This study aims to address the challenges of PL-based methods to improve SSMIS.

To address these challenges, a framework named Uncertainty-Aware Collaborative Learning with Mixed Images is proposed for SSMIS. Specifically, a pseudo-label uncertainty estimation method called Mixed Prototype guided Uncertainty Estimation (MPUE) is proposed, which estimates the uncertainty of pseudo-labels based on mixed prototypes of labeled and unlabeled data to reduce the impact of noise in pseudo-labels. Moreover, the proposed Collaborative Learning with Mixed Images (CLMI) leverages the advantages of perturbed and unperturbed images. CLMI learns from perturbed images to improve the ability to handle unknown data. Meanwhile, by enhancing the learning of unperturbed images, CLMI mitigates overfitting to unreliable predictions of perturbed images.

Experiments on the ACDC and PROMISE12 datasets with 10 % labeled data demonstrate the effectiveness of our framework, achieving Dice scores of 89.97 % and 81.20 %, respectively. The results significantly outperform state-of-the-art semi-supervised methods.

Our framework improves the accuracy of SSMIS. Extensive experiments on two public medical image datasets demonstrate the superiority of our framework. The code is available at: https://github.com/sjmei2004/Uncertainty-Aware-Collaborative-Learning-with-Mixed-Images.