A Probabilistic U-Net for Segmentation of Ambiguous Images

Abstract

Many real-world vision problems suffer from inherent ambiguities. In clinical applications for example, it might not be clear from a CT scan alone which particular region is cancer tissue. Therefore a group of graders typically produces a set of diverse but plausible segmentations. We consider the task of learning a distribution over segmentations given an input. To this end we propose a generative segmentation model based on a combination of a U-Net with a conditional variational autoencoder that is capable of efficiently producing an unlimited number of plausible hypotheses. We show on a lung abnormalities segmentation task and on a Cityscapes segmentation task that our model reproduces the possible segmentation variants as well as the frequencies with which they occur, doing so significantly better than published approaches. These models could have a high impact in real-world applications, such as being used as clinical decision-making algorithms accounting for multiple plausible semantic segmentation hypotheses to provide possible diagnoses and recommend further actions to resolve the present ambiguities.

arXiv commentLast update: added further details about the LIDC experiment. 11 pages for the main paper, 28 pages including appendix. 5 figures in the main paper, 18 figures in total, Advances in Neural Information Processing Systems (NeurIPS), 2018

References in corpus (11)

• Adam: A Method for Stochastic Optimization
• Auto-Encoding Variational Bayes
• What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?
• Semi-Supervised Learning with Deep Generative Models
• NIPS 2016 Tutorial: Generative Adversarial Networks
• Toward Multimodal Image-to-Image Translation
• The Cramer Distance as a Solution to Biased Wasserstein Gradients
• Why M Heads are Better than One: Training a Diverse Ensemble of Deep Networks
• Improving GANs Using Optimal Transport
• A note on the triangle inequality for the Jaccard distance
• Uncertainty Estimates and Multi-Hypotheses Networks for Optical Flow

Cited by in corpus (22)

• Uncertainty-Informed Deep Learning Models Enable High-Confidence Predictions for Digital Histopathology
• Effects of Image Size on Deep Learning
• QU-BraTS: MICCAI BraTS 2020 Challenge on Quantifying Uncertainty in Brain Tumor Segmentation - Analysis of Ranking Scores and Benchmarking Results
• BerDiff: Conditional Bernoulli Diffusion Model for Medical Image Segmentation
• Automated quantification of myocardial tissue characteristics from native T1 mapping using neural networks with Bayesian inference for uncertainty-based quality-control
• Learning Shape Priors for Robust Cardiac MR Segmentation from Multi-view Images
• Learning the Effect of Registration Hyperparameters with HyperMorph
• Automated Detection of Cribriform Growth Patterns in Prostate Histology Images
• Learning From Noisy Labels By Regularized Estimation Of Annotator Confusion
• Uncertainty-aware performance assessment of optical imaging modalities with invertible neural networks
• AlignShift: Bridging the Gap of Imaging Thickness in 3D Anisotropic Volumes
• Bayesian Neural Networks for Uncertainty Estimation of Imaging Biomarkers
• Out of distribution detection for intra-operative functional imaging
• Probabilistic Deep Learning for Instance Segmentation
• Effect of latent space distribution on the segmentation of images with multiple annotations
• BeSt-LeS: Benchmarking Stroke Lesion Segmentation using Deep Supervision
• Weakly- and Semi-Supervised Probabilistic Segmentation and Quantification of Ultrasound Needle-Reverberation Artifacts to Allow Better AI Understanding of Tissue Beneath Needles
• Variational Topic Inference for Chest X-Ray Report Generation
• Diffusion Based Ambiguous Image Segmentation
• Quantification of Predictive Uncertainty via Inference-Time Sampling
• Contrastive Learning with Temporal Correlated Medical Images: A Case Study using Lung Segmentation in Chest X-Rays
• Confidence Contours: Uncertainty-Aware Annotation for Medical Semantic Segmentation