Multi-task contrastive learning for semi-supervised medical image segmentation with multi-scale uncertainty estimation

Abstract

Objective. Automated medical image segmentation is vital for the prevention and treatment of disease. However, medical data commonly exhibit class imbalance in practical applications, which may lead to unclear boundaries of specific classes and make it difficult to effectively segment certain tail classes in the results of semi-supervised medical image segmentation.Approach. We propose a novel multi-task contrastive learning framework for semi-supervised medical image segmentation with multi-scale uncertainty estimation. Specifically, the framework includes a student-teacher model. We introduce global image-level contrastive learning in the encoder to address the class imbalance and local pixel-level contrastive learning in the decoder to achieve intra-class aggregation and inter-class separation. Furthermore, we propose a multi-scale uncertainty-aware consistency loss to reduce noise caused by pseudo-label bias.Main results. Experiments on three public datasets ACDC, LA and LiTs show that our method achieves higher segmentation performance compared with state-of-the-art semi-supervised segmentation methods.Significance. The multi-task contrastive learning in our method facilitates the negative impact of class imbalance and achieves better classification results. The multi-scale uncertainty estimation encourages consistent predictions for the same input under different perturbations, motivating the teacher model to generate high-quality pseudo-labels. Code is available athttps://github.com/msctransu/MCSSMU.git.

Keywords: contrastive learning; medical image segmentation; multi-task learning; semi-supervised learning; uncertainty estimation.