Uncertainty-weighted and relation-driven consistency training for semi-supervised head-and-neck tumor segmentation

Abstract

Accurately segmenting nasopharyngeal carcinoma (NPC), a type of head-and-neck tumor, in computed tomography (CT) images is essential for clinical treatment. However, due to the limited number of annotated samples, large shape variations, and boundary ambiguity, NPC segmentation remains a challenging task. Semi-supervised learning (SSL), which makes use of unlabeled data, has become a popular solution for NPC segmentation. In this paper, we propose a consistency training framework for semi-supervised NPC segmentation, which includes an uncertainty-weighted prediction consistency training (UPCT) strategy and a relation-driven consistency training (RCT) strategy. Our model is an encoder–decoder architecture consisting of a shared encoder, a main decoder, and several auxiliary decoders. We apply various perturbations to the shared encoder’s output to leverage the unlabeled data and enforce consistency between the predictions of the main and auxiliary decoders. To avoid being misled by unreliable outputs during training due to annotation scarcity, we introduce uncertainty estimation into our model. The UPCT strategy emphasizes predictions with high reliability and weakens unreliable predictions. Additionally, we propose the RCT strategy to explore intrinsic relations among samples in each mini-batch, encouraging the model to extract richer semantic information. Our extensive experiments on the in-house NPC dataset and the public 2017 ACDC dataset not only demonstrate the effectiveness of the UPCT and RCT strategies but also show the superiority and generalization of our method compared to the state-of-the-art methods.