Calibrated Uncertainty-Guided Multi-task Framework for Medical Image Segmentation

Yu Chen, Chunwei Wu, Shasha Liu, Guitao Cao

Abstract: Medical image segmentation is a crucial part of computer-aided diagnosis. Due to the enormous cost of labeling medical images, researchers have turned to exploring semi-supervised learning. However, the lack of supervisory information makes it difficult to accurately segment the fuzzy regions (e.g., complex edges or corners of organs). In this paper, we propose a novel method called Multi-task Consistency Segmentation Network based on Calibrated Uncertainty (CU-MCSNet). This model incorporates calibrated uncertainty to guide the network’s learning process. In addition, the model consists of two tasks: i) semantic segmentation as the primary task and ii) signed distance regression as the auxiliary task. To enhance the accuracy of edge segmentation, we propose the Edge Calibration Network for the primary task. This network integrates essential spatial and channel features, employing gradient complementation to hinder the accumulation of defective information and supply pertinent data to fuzzy regions. We also use the inter-task consistency loss to explore the underlying information of the images. In the multi-task domain, it is tough to balance each task manually, and we note that homoscedastic uncertainty focuses on inter-task variation. However, its numerical estimation may still be subject to bias. Therefore, we propose an adaptive loss-balancing strategy based on calibrated homoscedastic uncertainty. Extensive experiments show that our proposed method achieves state-of-the-art performance.