Diffusion Models for Implicit Image Segmentation Ensembles

Abstract

Diffusion models have shown impressive performance for generative modelling of images. In this paper, we present a novel semantic segmentation method based on diffusion models. By modifying the training and sampling scheme, we show that diffusion models can perform lesion segmentation of medical images. To generate an image specific segmentation, we train the model on the ground truth segmentation, and use the image as a prior during training and in every step during the sampling process. With the given stochastic sampling process, we can generate a distribution of segmentation masks. This property allows us to compute pixel-wise uncertainty maps of the segmentation, and allows an implicit ensemble of segmentations that increases the segmentation performance. We evaluate our method on the BRATS2020 dataset for brain tumor segmentation. Compared to state-of-the-art segmentation models, our approach yields good segmentation results and, additionally, detailed uncertainty maps.

SHAPR predicts 3D cell shapes from 2D microscopic images

“…1 for a visualization of the diffusion process during inference). Wolleb et al [22] used a U-Net model that in each step is trained to predict the probability density function f (x t ) from x t for all t ∈ {1, ..., T }, where x t−1 serves as the ground truth. Denote, with the model parameters θ, we can then write the reverse process p θ as…”

“…2). Following Wolleb et al [22] we concatenate the 2D image b in each forward pass through the model with x b , while noise is only added to the groundtruth x b . This leads to…”

“…We split the dataset into five folds with a 80%/20% train/test split, ensuring that each image is contained in the test set exactly once. We trained five separate models, one for each split for 85000 steps with T = 1000, and the hyperparameters provided by Wolleb et al [22]. We trained and evaluated the model on our computer cluster, containing Tesla V100 GPUs with 16GB memory.…”

“…Denoising diffusion probabilistic models were further improved in their loss function by architecture changes and by classifier guidance during sampling, leading to improved image quality of the predictions [12,3]. While most diffusion models are applied in the natural image domain, Wolleb et al have used them for segmentation of MRI images [22] and anomaly detection [21] in multimodal brain images, showing their applicability in the medical domain for segmentation of 2D MRI images. We propose a diffusion model, hereafter called DIffusion based Shape PRediction, DISPR for predicting 3D single cell shapes that are realistic reconstructions of 2D microscopy images.…”

“…Note that the zoom level in the visualizations is not uniform. fusion models are typically constrained to match the input domain by adding a prior, such as a text encoding [14] or a segmentation mask of a medical image [22]. Here, we apply a diffusion model to the inverse problem of generating three-dimensional (3D) shapes from two-dimensional (2D) microscopy images, using the 2D image as a prior.…”