A unified uncertainty network for tumor segmentation using uncertainty cross entropy loss and prototype similarity

摘要

Uncertainty estimation and out-of-distribution (OOD) detection are topics of practical significance in deep convolutional neural network-based tumor segmentation tasks. We propose a unified uncertainty segmentation network to handle these two tasks with a single network. The uncertainty cross entropy loss is proposed to guide the network to directly output the prediction uncertainty of each pixel instead of executing several times in the prediction phase. We extend the uncertainty to the case level to address OOD sample detection problems. Case-level uncertainty is estimated based on prototype similarity. We perform pixel-level uncertainty experiments and OOD detection experiments on four datasets. The experimental results show that our proposed method is more suitable for uncertainty estimation in tumor segmentation than existing methods. Our proposed method only needs to modify the network output and loss function and does not need to execute the network multiple times when estimating uncertainty. Moreover, our method shows improved performance in tumor segmentation.