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

arXiv commentAccepted 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 (20)

• Adam: A Method for Stochastic Optimization
• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Auto-Encoding Variational Bayes
• Categorical Reparameterization with Gumbel-Softmax
• What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?
• An Introduction to Variational Autoencoders
• Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning
• Stochastic Backpropagation and Approximate Inference in Deep Generative Models
• Variational Inference with Normalizing Flows
• The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
• Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles
• Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks
• Variational Dropout and the Local Reparameterization Trick
• Standardized Assessment of Automatic Segmentation of White Matter Hyperintensities and Results of the WMH Segmentation Challenge
• Bayesian Uncertainty Estimation for Batch Normalized Deep Networks
• A Probabilistic U-Net for Segmentation of Ambiguous Images
• Cascaded V-Net using ROI masks for brain tumor segmentation
• Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty
• 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