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

TL;DR

This paper introduces a probabilistic image segmentation model that integrates shape geometry into deep neural networks, enabling the generation of multiple plausible, geometry-aware segmentations that reflect model and input uncertainties.

Contribution

It combines the Probabilistic U-Net with Kendall Shape Variational Auto-Encoder to incorporate shape geometry into probabilistic segmentation, improving spatial coherence.

Findings

Produces more robust, geometry-aware segmentations.

Generates multiple plausible segmentations reflecting uncertainties.

Enhances spatial coherence in segmented regions.

Abstract

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 \cite{kohl2018probabilistic} to generate probabilistic segmentations, i.e.\! multiple likely segmentations for an input image. Our model also adopts the Kendall Shape Variational Auto-Encoder of \cite{vadgama2023kendall} 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.