# Probabilistic U-Net with Kendall Shape Spaces for Geometry-Aware Segmentations of Images

arXiv: https://arxiv.org/abs/2410.14017 · HTML: https://mlsanity.com/paper/2410.14017/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 32 references that are in the map.
- First posted on arXiv 17 Oct 2024.

## Builds on (most cited first, 18 of 18)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Image Segmentation Using Deep Learning: A Survey](https://mlsanity.com/paper/2001.05566/) (2020)
- [General E(2)-Equivariant Steerable CNNs](https://mlsanity.com/paper/1911.08251/) (2019)
- [Steerable CNNs](https://mlsanity.com/paper/1612.08498/) (2016)
- [A Probabilistic U-Net for Segmentation of Ambiguous Images](https://mlsanity.com/paper/1806.05034/) (2018)
- [Bayesian SegNet: Model Uncertainty in Deep Convolutional Encoder-Decoder Architectures for Scene Understanding](https://mlsanity.com/paper/1511.02680/) (2015)
- [Learning Steerable Filters for Rotation Equivariant CNNs](https://mlsanity.com/paper/1711.07289/) (2017)
- [Hyperspherical Variational Auto-Encoders](https://mlsanity.com/paper/1804.00891/) (2018)
- [Why M Heads are Better than One: Training a Diverse Ensemble of Deep Networks](https://mlsanity.com/paper/1511.06314/) (2015)
- [Uncertainty Estimates and Multi-Hypotheses Networks for Optical Flow](https://mlsanity.com/paper/1802.07095/) (2018)
- [Learning in an Uncertain World: Representing Ambiguity Through Multiple Hypotheses](https://mlsanity.com/paper/1612.00197/) (2016)
- [Understanding Deep Learning Techniques for Image Segmentation](https://mlsanity.com/paper/1907.06119/) (2019)
- [Learning Probabilistic Topological Representations Using Discrete Morse Theory](https://mlsanity.com/paper/2206.01742/) (2022)
- [A Note on the Kullback-Leibler Divergence for the von Mises-Fisher distribution](https://mlsanity.com/paper/1502.07104/) (2015)
- [Joint M-Best-Diverse Labelings as a Parametric Submodular Minimization](https://mlsanity.com/paper/1606.07015/) (2016)
- [A Sheaf-Theoretic Construction of Shape Space](https://mlsanity.com/paper/2204.09020/) (2022)

## Built on by (most cited first, 1 of 1)

- [Uncertainty-Guided Selective Adaptation Enables Cross-Platform Predictive Fluorescence Microscopy](https://mlsanity.com/paper/2511.12006/) (2025)

## Abstract (arXiv)

One of the fundamental problems in computer vision is image segmentation, the task of detecting distinct regions or objects in given images. Deep Neural Networks (DNN) have been shown to be very effective in segmenting challenging images, producing convincing segmentations. There is further need for probabilistic DNNs that can reflect the uncertainties from the input images and the models into the computed segmentations, in other words, new DNNs that can generate multiple plausible segmentations and their distributions depending on the input or the model uncertainties. While there are existing probabilistic segmentation models, many of them do not take into account the geometry or shape underlying the segmented regions. In this paper, we propose a probabilistic image segmentation model that can incorporate the geometry of a segmentation. Our proposed model builds on the Probabilistic U-Net of citekohl2018probabilistic to generate probabilistic segmentations, i.e.! multiple likely segmentations for an input image. Our model also adopts the Kendall Shape Variational Auto-Encoder of citevadgama2023kendall to encode a Kendall shape space in the latent variable layers of the prior and posterior networks of the Probabilistic U-Net. Incorporating the shape space in this manner leads to a more robust segmentation with spatially coherent regions, respecting the underlying geometry in the input images.