Learning to segment subcortical structures from noisy annotations with a novel uncertainty-reliability aware learning framework

Abstract

Accurate segmentation of subcortical structures is an important task in quantitative brain image analysis. Convolutional neural networks (CNNs) have achieved remarkable results in medical image segmentation. However, due to the difficulty of acquiring high-quality annotations of brain subcortical structures, learning segmentation networks using noisy annotations is an inevitable topic. A common practice is to select images or pixels with reliable annotations for training, which usually may not make full use of the information from the training samples, thus affecting the performance of the learned segmentation model. To address the above problem, in this work, we propose a novel robust learning method and denote it as uncertainty-reliability awareness learning (URAL), which can make sufficient use of all training pixels. At each training iteration, the proposed method first selects training pixels with reliable annotations from the set of pixels with uncertain network prediction, by utilizing a small clean validation set following a meta-learning paradigm. Meanwhile, we propose the online prototypical soft label correction (PSLC) method to estimate the pseudo-labels of label-unreliable pixels. Then, the segmentation loss of label-reliable pixels and the semi-supervised segmentation loss of label-unreliable pixels are used to calibrate the total segmentation loss. Finally, we propose a category-wise contrastive regularization to learn compact feature representations of all uncertain training pixels. Comprehensive experiments are performed on two publicly available brain MRI datasets. The proposed method achieves the best Dice scores and MHD values on both datasets compared to several recent state-of-the-art methods under all label noise settings. Our code is available at https://github.com/neulxlx/URAL.