PixelSeg: Pixel-by-Pixel Stochastic Semantic Segmentation for Ambiguous Medical Images

Semantic segmentation tasks often have multiple output hypotheses for a single input image. Particularly in medical images, these ambiguities arise from unclear object boundaries or differences in physicians' annotation. Learning the distribution of annotations and automatically giving multiple plausible predictions is useful to assist physicians in their decision-making. In this paper, we propose a semantic segmentation framework, PixelSeg, for modelling aleatoric uncertainty in segmentation maps and generating multiple plausible hypotheses. Unlike existing works, PixelSeg accomplishes the semantic segmentation task by sampling the segmentation maps pixel by pixel, which is achieved by the PixelCNN layers used to capture the conditional distribution between pixels. We propose (1) a hierarchical architecture to model high-resolution segmentation maps more flexibly, (2) a fast autoregressive sampling algorithm to improve sampling efficiency by 96.2, and (3) a resampling module to further improve predictions' quality and diversity. In addition, we demonstrate the great advantages of PixelSeg in the novel area of interactive uncertainty segmentation, which is beyond the capabilities of existing models. Extensive experiments and state-of-the-art results on the LIDC-IDRI and BraTS 2017 datasets demonstrate the effectiveness of our proposed model.