Generalized Probabilistic U-Net for medical image segementation

TL;DR

This paper introduces the Generalized Probabilistic U-Net, enhancing uncertainty modeling in medical image segmentation by using more flexible Gaussian-based latent spaces, leading to improved diversity and accuracy of predictions.

Contribution

It extends the Probabilistic U-Net with a more general Gaussian latent space, improving uncertainty approximation and segmentation performance.

Findings

Mixture of Gaussians improves GED metric.

Choice of latent distribution affects sample diversity.

Implementation available online.

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 uncertainty in the reference segmentations using the LIDC-IDRI dataset. We show that the choice of distribution affects the sample diversity of the predictions and their overlap with respect to the reference segmentations. For the LIDC-IDRI dataset, we show that using a mixture of Gaussians results in a statistically significant improvement in the generalized energy distance (GED) metric with respect to the standard Probabilistic U-Net. We have made our implementation available at https://github.com/ishaanb92/GeneralizedProbabilisticUNet