Semi-supervised NPC segmentation with uncertainty and attention guided consistency

摘要

Segmentation of nasopharyngeal carcinoma (NPC) from computed tomography (CT) image is conducive to the clinical healthcare. Nevertheless, due to the large shape variations, boundary ambiguity, as well as the limited available annotations, NPC segmentation remains to be a challenging task. In this paper, we propose a two-stage semi-supervised segmentation framework for automatic NPC segmentation, which includes a region of interest (ROI) cropping stage and a semi-supervised segmentation stage. Specifically, considering the large individual variability of NPC tumors, we first employ a coarse-Res-Unet (CRU) to extract the rough target areas from the CT images and thus obtain the cropped ROI images. Then, both the entire CT images and the corresponding ROI images are input to a self-attention embedded semi-supervised mean teacher (SSMT) model to generate the ROI-focused segmentation results. Here, to relieve the potential misdirection from the teacher model caused by annotation scarcity, we introduce the uncertainty estimation to guide the student model to gradually learn the reliable predictions from the teacher model. Meanwhile, to fully explore the inherent semantic information of unlabeled data, we also encourage the attention maps from these two models to be consistent at feature level. Furthermore, we design a refinement procedure and reuse the ROI attention maps generated by the well-trained SSMT to retrain the first stage, improving the quality of ROI images. The updated ROI images are further leveraged to refine SSMT to predict the final segmentation results. Note that the uncertainty estimation and the attention maps reveal the confidence and attention of the model for the intermediate features respectively, which can provide explainable evaluation to the quality of segmentation results. Experimental results on an in-house NPC dataset and a public 2017 ACDC dataset demonstrate that our method is superior to other semi-supervised segmentation methods and also has good generalization ability.