Uncertainty teacher with dense focal loss for semi-supervised medical image segmentation

Abstract

In medical scenarios, obtaining pixel-level annotations for medical images is expensive and time-consuming, even if considering its importance for automating segmentation tasks. Due to the scarcity of labels in the training phase, semi-supervised methods are widely applied for various medical tasks. To better utilize the unlabeled data, several works have explored the method of uncertainty estimation and exhibited huge success. Despite their impressive performance, we believe that the underlying information of the unlabeled data has been largely unexplored. Meanwhile, there is an extreme foreground–background class imbalance during the training phase of semantic segmentation, which may cause a vast number of easily classified samples to overwhelm the loss during training and lead to a model collapse. In this paper, we proposed uncertainty teacher with dense focal loss, a method that can take good advantage of unlabeled data simultaneously and address the class imbalance problem, based on Deep Co-Training. On one hand, the uncertainty teacher framework is presented to better utilize the unlabeled data by introducing a novel method to regularize uncertainty in the right direction, and the uncertainty is estimated by Monte Carlo Sampling. On the other hand, the dense focal loss is proposed to help solve the class imbalance problem between different classes of samples in medical image segmentation and effectively convert the multi-variate entropy into a multiple binary entropy. We implemented our method on three challenging public medical datasets and experimental results have shown desirable improvements to state-of-the-art.