Effect of latent space distribution on the segmentation of images with multiple annotations

Abstract

We propose the Generalized Probabilistic U-Net, which extends the Probabilistic U-Net by allowing more general forms of the Gaussian distribution as the latent space distribution that can better approximate the uncertainty in the reference segmentations. We study the effect the choice of latent space distribution has on capturing the variation in the reference segmentations for lung tumors and white matter hyperintensities in the brain. We show that the choice of distribution affects the sample diversity of the predictions and their overlap with respect to the reference segmentations. We have made our implementation available at https://github.com/ishaanb92/GeneralizedProbabilisticUNet

Accepted for publication at the Journal of Machine Learning for Biomedical Imaging (MELBA) https://melba-journal.org/2023:005. arXiv admin note: text overlap with arXiv:2207.12872

References in corpus (6)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
• Standardized Assessment of Automatic Segmentation of White Matter Hyperintensities and Results of the WMH Segmentation Challenge
• Cascaded V-Net using ROI masks for brain tumor segmentation
• RDCNet: Instance segmentation with a minimalist recurrent residual network

Cited by in corpus (2)

• Investigating and Improving Latent Density Segmentation Models for Aleatoric Uncertainty Quantification in Medical Imaging
• Diffusion Based Ambiguous Image Segmentation