Effect of latent space distribution on the segmentation of images with multiple annotations
This work addresses uncertainty modeling in medical image segmentation for clinicians, but it is incremental as it builds directly on the Probabilistic U-Net.
The authors tackled the problem of capturing uncertainty in medical image segmentation by extending the Probabilistic U-Net to use more general Gaussian distributions in the latent space, showing that this choice affects prediction diversity and overlap with reference segmentations for lung tumors and white matter hyperintensities.
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