A DIFFUSION MODEL PREDICTS 3D SHAPES FROM 2D MICROSCOPY IMAGES

1,3,⋆ 1,2 1,2,† 4,† 1,†

Dominik J. E. Waibel Ernst Röell Bastian Rieck Raja Giryes Carsten Marr

1
Institute of AI for Health, Helmholtz Munich – German Research Centre for Environmental Health, Neuherberg, Germany
2
TUM School of Computation, Information and Technology, Technical University of Munich, Munich Germany
3
TUM School of Life Sciences, Technical University of Munich, Munich, Germany
4
Faculty of Engineering, Tel Aviv University, Tel Aviv, Israel
⋆
Current affiliation: AstraZeneca Computational Pathology, Oncology R&D, Munich, Germany
†
These authors share corresponding authorship
arXiv:2208.14125v3 [[Link]] 14 Mar 2023

ABSTRACT This is due to the fact that the sequential imaging and stacking of
slices of 2D images is time-consuming and toxic for cells. However,
Diffusion models are a special type of generative model, capable
the benefits of 3D microscopy are staggering: assessing morpholog-
of synthesising new data from a learnt distribution. We introduce
ical information of individual cells—such as their volume, shape or
DISPR, a diffusion-based model for solving the inverse problem of
surface—via 3D microscopy promises new insights into blood dis-
three-dimensional (3D) cell shape prediction from two-dimensional
orders such as sickle cell anemia [1]. Imaging a large number of
(2D) single cell microscopy images. Using the 2D microscopy im-
cells at high resolution in 3D is rather costly. Therefore, researchers
age as a prior, DISPR is conditioned to predict realistic 3D shape re-
have to find ways how they can optimally balance throughput and
constructions. To showcase the applicability of DISPR as a data aug-
resolution of 3D microscopy data.
mentation tool in a feature-based single cell classification task, we
extract morphological features from the red blood cells grouped into
six highly imbalanced classes. Adding features from the DISPR pre-
dictions to the three minority classes improved the macro F1 score
from F 1macro = 55.2 ± 4.6% to F 1macro = 72.2 ± 4.9%. We thus
demonstrate that diffusion models can be successfully applied to in-
verse biomedical problems, and that they learn to reconstruct 3D
shapes with realistic morphological features from 2D microscopy
images.

1. INTRODUCTION

Diffusion models are a class of generative models, showing supe- Fig. 2. DISPR is trained to denoise the 3D volume xb,t containing
rior performance as compared to other generative models in creating added stochastic Gaussian noise, to obtain xb,t−1 , thus reversing the
realistic images when trained on natural image datasets. We apply noising step q. In each forward pass pθ , we constrain our 3D model
a diffusion model to the reconstruction of three-dimensional single with one 2D image b containing a fluorecent image and a mask. Dur-
cell shapes (3D) from two-dimensional (2D) microscopy images, us- ing inference, the forward pass pθ is repeated T times to obtain the
ing the 2D image as a prior. prediction xb,0 (see section 2).

A second challenge in the medical domain, especially in 3D

microscopy, is the manual labor involved in labeling, segmenting
and annotating data samples. This dependence makes it both time
consuming and expensive. Moreover, datasets are most likely im-
balanced, further impeding the application of deep learning meth-
ods. Data augmentation techniques generate synthetic training sam-
ples from the training dataset, adding to the total volume of training
data available for deep learning models. The distributional nature of
our model lends itself particularly well for data augmentation. We
demonstrate this in a red blood cell classification task, where we
augment the training set with synthetic samples and train a classifier.
Fig. 1. Diffusion steps of one forward pass through DISPR visual- Existing approaches [2, 3] aim to predict a single, determinis-
ized during inference between step t = 999 with the largest amount tic reconstruction of a biological structure for each single 2D in-
of Gaussian noise and the DISPR prediction of t = 0, for a stroma- put. Others focus on predicting the shapes of natural objects such
tocyte cell (see section 2). Note that the zoom level in the visualiza- as air planes, cars, and furniture from photographs, creating either
tions is not uniform. meshes [4, 5], voxel volumes [6, 7], or point clouds [8]. In the
biomedical domain, SHAPR, a deterministic autoencoder, was de-
Our work is motivated by the observation that imaging speed is veloped by Waibel et al. to predict the shape and morphology of in-
a central limitation for microscopic confocal imaging of single cells. dividual mammalian cells from 2D microscopy images [2]. SHAPR
 2. METHODS

We first provide a brief overview of diffusion models in general be-

fore detailing the specific architectural choices of DISPR, our model
for predicting 3D cell shapes from 2D microscopy images.
A diffusion model contains two processes, a forward diffusion
process and a reconstruction (or denoising) process. In the forward
diffusion process, Gaussian noise is successively added to an input
image in a predefined number of steps, generating a sequence of im-
ages with increasing noise. This sequence will in the limit approach
a normal distribution with unit variance, losing all information. The
sequences of noisy images will then serve as the training data for
the reconstruction process. The diffusion model is trained to recover
the original image given the sequence of noisy images, in essence
denoising them.
To describe the forward diffusion process, let x0 denote the orig-
inal 3D image. We denote the normal distribution with mean µ and
variance β applied to xt by N (xt ; µ, β). The forward diffusion of
xt−1 to xt is recursively given by
√
q(xt ∣xt−1 ) = N (xt ; 1 − βt xt−1 , βt I), (1)

where I denotes the identity matrix. Note that we allow the variance
βt to increase with each step [20, 10]. The forward diffusion process
is repeated for a fixed, predefined number of steps T . In our case,
in line with the literature, we set T = 1000. Repeated application
of Eq. (1) to the original image x0 and setting αt = 1 − βt and
t
ᾱt = ∏i=1 αi yields
√
q(xt ∣x0 ) = N (xt ; ᾱt x0 , (1 − ᾱt )I). (2)
Fig. 3. Red blood cells exhibit similar morphological features be-
tween groundtruth (3rd column, yellow background) and DISPR Hence xt can be written in terms of x0 [20, 10] as
predictions (white background). A 2D microscopy image (with a √ √
xt = ᾱt x0 + 1 − ᾱt  with  ∼ N (0, I). (3)
fluorescence and segmented channel, dark background) is used as
the model’s input to predict the 3D shapes. The diffusion model is trained to reverse the forward process by re-
moving noise with each forward pass p through the model, trying to
predict xt−1 from the input xt (see Fig. 2 for a visualization of the
has recently been refined with a topological loss function [3]. Both diffusion process during inference). Wolleb et al. [13] used a U-Net
methods are based on an autoencoder that predicts one shape for model that in each step is trained to predict the probability density
each input image. By contrast, our method can be used to predict a function f (xt ) from xt for all t ∈ {1, ..., T }, where xt−1 serves as
distribution of 3D shape reconstructions for each input image. the ground truth. With the model parameters denoted by θ, we can
then write the reverse process pθ as
Diffusion models have gained much attention due to their as-
tonishing performance in generating realistically-looking images [9, pθ (xt−1 ∣xt ) = N (xt−1 ; µθ (xt , t), Σθ (xt , t)). (4)
10]. Denoising diffusion implicit models [11] have substantially im-
proved the noise sampling scheme by skipping multiple noise sam- Ho et al. [20] derive the formula for the forward pass of the model
pling steps. Denoising diffusion probabilistic models were further as
1 1 − αt
improved in their loss function by architecture changes and by clas- xt−1 = √ (xt √ θ (xt , t)) + σt z, (5)
sifier guidance during sampling, leading to improved image quality αt 1 − ᾱt
of the predictions [10, 12]. While most diffusion models are ap- with σt being the variance scheme the model can learn [10]. Com-
plied in the natural image domain, Wolleb et al. have used them ponent z in equation 5 reflects the√stochastic√sampling process. The
for segmentation of MRI images [13] and anomaly detection [14] model is trained with input xt = ᾱt x0 + 1 − ᾱt  to subtract the
in multimodal brain images, showing their applicability in the med- noise scheme θ (xt , t) from xt according to equation 5.
ical domain for segmentation of 2D MRI images. Diffusion models
have been used of 2D MRI and CT image synthesis [15, 16, 17], to 2.1. DISPR
synthesise longitudinal MRI images [18] and 4D MRI images [19].
We go beyond previous work and propose a diffusion model, We use 2D images to constrain the 3D image generation. In each
2
1
hereafter called DIffusion based Shape PRediction, DISPR for pre- training and evaluation step the 2D image b ∈ R is concatenated to
3
dicting single cell shapes that are realistic 3D reconstructions from the noisy segmentation mask xb,t ∈ R . The groundtruth volume
3
2D microscopy images. for each input image is xb,0 ∈ R . This leads to

1 1 − αt
1
See [Link] xb,t−1 = √ (xb,t √ θ (xb,t ⊕ b, t)) + σt z. (6)
αt 1 − ᾱt
 Relative
Model Median µ±σ
error
Cylinder fit 0.32 0.34 ± 0.22
Volume Ellipsoid fit 0.40 0.37 ± 0.23
SHAPR 0.15 0.20 ± 0.18
Topo SHAPR 0.26 0.29 ± 0.27
DISPR (ours) 0.15 0.20 ± 0.20
Cylinder fit 0.30 0.30 ± 0.19
Surface Ellipsoid fit 0.41 0.42 ± 0.19
area SHAPR 0.15 0.16 ± 0.11
Topo SHAPR 0.14 0.18 ± 0.16
DISPR (ours) 0.11 0.14 ± 0.15
Cylinder fit 0.26 0.29 ± 0.17
Surface
Ellipsoid fit 0.48 0.49 ± 0.19
roughness
SHAPR 0.32 0.31 ± 0.16
Topo SHAPR 0.25 0.29 ± 0.29
DISPR (ours) 0.15 0.23 ± 0.27
Cylinder fit 4.79 4.76 ± 1.36
Curvature Ellipsoid fit 0.38 0.39 ± 0.17
SHAPR 0.30 0.32 ± 0.21
Topo SHAPR 0.28 0.34 ± 0.24
DISPR (ours) 0.17 0.19 ± 0.12

Fig. 4. With the predictions of DISPR we obtain a lower relative

volume error as compared to the extrapolations of two naive models, Table 1. Median, mean (µ), and standard deviation (σ) of relative
the cylinder and ellipsoid fit, as well as predictions of SHAPR [2] errors (lower values are better; best result is marked in bold).
and the predictions of SHAPR using the topological loss [3]. DISPR
also outperforms the other models with respect to the surface area
error, surface roughness error, and relative surface curvature error.
Note that the relative errors of the predictions of DISPR contain five We trained five separate models, one for each split for 85000 steps
times as many datapoints than the others. with T = 1000, and the hyper-parameters provided by Wolleb et
al. [13]. We trained and evaluated the model on a compute clus-
ter containing Tesla V100 GPUs with 16GB memory. Due to GPU
During inference, a noisy image xb,T , is passed T times through memory constraints we had to run the models with a batch-size of
−4
the model (Fig. 3). Because of the stochasticity of xb,T the diffusion one. We used a learning rate of 10 , with no weight decay, a linear
model fθ (xT,b ) = pθ (...(pθ (xT,b )) predicts different outputs x0 af- noise scheduler, and the Adam optimizer. The number of channels
ter T = 1000 forward passes through the model, which all represent in the first layer were 32, and the layers of our multihead attention
realistic reconstructions of the 2D image constraint b (see Fig. 1). layers had a resolution of 16. We used a uniform sample scheduler
for the diffusion steps. Each model was trained for 8500 steps which
2.2. Dataset took roughly 60 hours, summing up to 300 hours of training time for
five models. During inference the model took on average 4:45 min-
We use a publicly-available red blood cell dataset published by utes for each prediction, summing up to approximately 330 hours for
Simionato et al. [21]. It consists of 825 3D images of red blood cells 4125 predictions, as we predict five 3D cell shapes for each 2D input
recorded with a confocal microscope. All cells are available as seg- image.
mented 3D images [2, 22]. Each image is contained in a 64×64×64 We evaluate each of the five trained diffusion models on ev-
voxel grid. Each cell is assigned to one of the following six classes, ery image b in the test set five times to obtain five shape predic-
with n beeing the number of images in each class: (i) SDE shapes tions. As each model is tested on 165 2D images, this adds up to
(n = 602), (ii) cell clusters (n = 69), (iii) multilobates (n = 12), 165 ⋅ 25 = 4125 predicted cell shapes. Due to the stochasticity of
(iv) keratocytes (n = 31), (v) knizocytes (n = 23), or (vi) acantho- the diffusion model, variations between the predicted shapes occur
cytes (n = 88). Note that Spherocytes, stomatocytes, discocytes, and each predicted cell is unique (see Fig. 3). To quantify the per-
and echinocytes are combined into the SDE shapes class, which formance of our model, we compare biologically relevant features
is characterized by the stomatocyte–discocyte–echinocyte trans- such as the volume, surface area, surface roughness, and convexity
formation [21]. The 2D images contained in the provided dataset to the groundtruth to reconstructions obtained from SHAPR [2] or its
have been extracted from the central slide of each 3D image and topology-aware variant [3]; to date, these models constitute the state-
segmented by thresholding, thus resulting in a (64,64) pixel size. of-the-art for this shape reconstruction task. As additional compar-
isons, we use two naive baseline models, a cylinder and ellipsoid fit,
3. EXPERIMENTS where we extrapolate the 3D shape from the 2D segmentation mask,
similar to Waibel et al. [2]. We calculate the mean of the major and
We split the dataset into five folds with a 80%/20% train/test split, the minor axis of an ellipse fitted to the 2D segmentation outline and
ensuring that each image is contained in the test set exactly once. use this as the cylinder’s extent in the third dimension and as the
 the test set exactly once. The random forest models are trained with
1000 estimators and a depth of 10.
To oversample the three smallest classes, we predict five 3D cell
shapes for each 2D image of the three smallest classes contained in
the dataset (multilobates, keratocytes, and knizocytes) (see Fig. 5a).
In a second step, we add features extracted from DISPR’s predictions
to the respective training sets, if the same cell was already contained
in the training dataset. Thereby we ensure that no overlap between
the train and test set occurs. By enriching the groundtruth features
with features obtained from the predictions of DISPR, the classifi-
cation performance significantly increased (F 1macro = 72.2 ± 4.9%,
F 1weighted = 88.6 ± 1.1%, mean±sdev, n = 5 cross-validation runs)
as compared to using groundtruth microscopy image features only
(F 1macro = 55.2 ± 4.6%, F 1weighted = 85.7 ± 2.1% (see Fig 5b)) in
a fivefold cross-validation (see Fig. 5b). The largest improvements
were found for those classes to which we added features, especially
for the keratocytes, the number of correctly classified cells increased
from 14 to 24 and for the multilobate cells from 1 to 8 correct clas-
sifications (see Fig 5c).

5. DISCUSSION

We demonstrate that DISPR reconstructs 3D shapes from 2D mi-

croscopy images by comparing relevant morphological features.
Class-specific morphological feature distributions from the ground
truth dataset are captured in the model distribution and suggests that
our approach is capable of data augmentation.
The lower relative surface roughness and curvature error further
suggest improvements in predicting morphological properties com-
pared to other state of the art models. This also gives credence to our
visual impression, that our model is able to predict 3D shapes that
are more realistic. Additionally, we show that the morphological fea-
tures do not only appear more realistic, but DISPR can be utilized for
Fig. 5. Using DISPR to oversample training data improves a feature- generating training data, used in a classification task, leading to sig-
based random forest classification of single red blood cells. (a) For nificant improvements especially on minority classes, emphasising
the training dataset, containing 128 manually extracted morpholog- the utility of DISPR for generating data. DISPR, is not restricted to
ical features of the 3D groundtruth, we enhance the three smallest single cells, and may be applied to other reconstruction tasks as well.
classes (multilobate, keratocyte, knizocyte) with features extracted We acknowledge that DISPR’s inference is magnitudes less effi-
from the DISPR’s predictions, reducing class imbalance. (b) We im- cient as compared to autoencoder models, such as SHAPR [2, 3] for
prove the macro F1 score to F 1macro = 72.2 ± 4.9%, from F 1macro 3D shape prediction, where each prediction takes a few seconds, as
= 55.2 ± 4.6%, with the largest improvements obtained for minority compared to 4:45minutes using DISPR.
classes (c). One future direction could be to consider geometrical and topo-
logical information in the model’s regularization to improve train-
ing efficiency and model performance. Exploiting the distributional
third axis for the ellipsoid fit, respectively. nature of predicted shapes could lead to additional studies of their
structure; cells whose reconstructions are highly ambiguous could
4. RESULTS be detected by a multi-modal shape reconstruction. Finally, to re-
duce computational cost during inference, which is one of the main
In our experiments we find high visual morphological consis- drawbacks of diffusion models, super-resolution models paired with
tency between the five predictions and also a high similarity to diffusion models [23] might prove useful in 3D applications. In
the groundtruth 3D shape. For individual features, such as the vol- the biomedical domain denoising images is a recurring task, which
ume, surface area, surface roughness and curvature our approach might be solvable with diffusion models, as their training objective
outperforms the four competing approaches [2, 3] (see Fig. 4 and itself is the denoising of images.
Table 1) in terms of relative error to the groundtruth.
To test if DISPR based training data generation improves sin- 5.1. Data and code availability
gle cell classification, we extract 128 morphological features from
each single cell, similar to Waibel et al. [2]. Using random forest The dataset is available at [22]: doi: 10.5281/zenodo.7031924
classifiers, we predict the respective class of each cell based on the The code is available at: [Link]
extracted morphological features. Five random forest models are
trained to classify each cell in one of the six classes suggested by 5.2. Compliance with Ethical Standards
Simionato et al. [21] in a round-robin fashion, splitting the dataset
No ethical approval was required for this study.
with a 80%/20% train/test split, ensuring each image is contained in
5.3. Acknowledgements [10] Alexander Quinn Nichol and Prafulla Dhariwal, “Improved de-
noising diffusion probabilistic models,” in Proceedings of the
We thank Julia Wolleb (Basel) for the inspiration, Melanie Schulz 38th International Conference on Machine Learning, Marina
and Daniel Lang (Munich), and Erez Yosef (Tel Aviv) for discussing Meila and Tong Zhang, Eds. 2021, vol. 139 of Proceedings of
ideas and feedback to this manuscript. We also gratefully acknowl- Machine Learning Research, pp. 8162–8171, PMLR.
edge the use of Helmholtz Munich High-Performance Computing
Cluster for model training. Using the ML CO2 calculator [24] we [11] Jiaming Song, Chenlin Meng, and Stefano Ermon, “Denoising
estimate 81.65kg of CO2 emitted, considering only experiments. diffusion implicit models,” arXiv preprint arXiv:2010.02502,
2020.
5.4. Funding [12] Dhariwal and Nichol, “Diffusion models beat gans on image
synthesis,” Adv. Neural Inf. Process. Syst.
Carsten Marr received funding from the European Research Coun- [13] Julia Wolleb, Robin Sandkühler, Florentin Bieder, Philippe
cil (ERC) under the European Union’s Horizon 2020 Research and Valmaggia, and Philippe C Cattin, “Diffusion models for
Innovation Programme (Grant Agreement 866411). implicit image segmentation ensembles,” arXiv preprint
arXiv:2112.03145, 2021.
6. REFERENCES [14] Julia Wolleb, Florentin Bieder, Robin Sandkühler, and
Philippe C Cattin, “Diffusion models for medical anomaly de-
[1] Monica Diez-Silva, Ming Dao, Jongyoon Han, Chwee-Teck tection,” arXiv preprint arXiv:2203.04306, 2022.
Lim, and Subra Suresh, “Shape and biomechanical character- [15] Firas Khader, Gustav Mueller-Franzes, Soroosh Tayebi
istics of human red blood cells in health and disease,” MRS Arasteh, Tianyu Han, Christoph Haarburger, Maximilian
bulletin, vol. 35, no. 5, pp. 382–388, 2010. Schulze-Hagen, Philipp Schad, Sandy Engelhardt, Bettina
[2] Dominik J.E. Waibel, Niklas Kiermeyer, Scott Atwell, Ario Baessler, Sebastian Foersch, et al., “Medical diffusion–
Sadafi, Matthias Meier, and Carsten Marr, “Shapr predicts 3d denoising diffusion probabilistic models for 3d medical image
cell shapes from 2d microscopic images,” iScience, vol. 25, no. generation,” arXiv preprint arXiv:2211.03364, 2022.
11, pp. 105298, 2022. [16] Muzaffer Özbey, Salman UH Dar, Hasan A Bedel, Onat Dal-
[3] Dominik J. E. Waibel, Scott Atwell, Matthias Meier, Carsten maz, Şaban Özturk, Alper Güngör, and Tolga Çukur, “Unsu-
Marr, and Bastian Rieck, “Capturing shape information pervised medical image translation with adversarial diffusion
with multi-scale topological loss terms for 3d reconstruction,” models,” arXiv preprint arXiv:2207.08208, 2022.
in Medical Image Computing and Computer Assisted Inter- [17] Hazrat Ali, Shafaq Murad, and Zubair Shah, “Spot the fake
vention – MICCAI 2022, Linwei Wang, Qi Dou, P. Thomas lungs: Generating synthetic medical images using neural dif-
Fletcher, Stefanie Speidel, and Shuo Li, Eds., Cham, 2022, pp. fusion models,” arXiv preprint arXiv:2211.00902, 2022.
150–159, Springer Nature Switzerland. [18] Jee Seok Yoon, Chenghao Zhang, Heung-Il Suk, Jia Guo,
[4] Georgia Gkioxari, Jitendra Malik, and Justin Johnson, “Mesh and Xiaoxiao Li, “Sadm: Sequence-aware diffusion model
R-CNN,” in Proceedings of the IEEE/CVF International Con- for longitudinal medical image generation,” arXiv preprint
ference on Computer Vision (ICCV), 2019. arXiv:2212.08228, 2022.
[5] Nanyang Wang, Yinda Zhang, Zhuwen Li, Yanwei Fu, Wei [19] Boah Kim and Jong Chul Ye, “Diffusion deformable model
Liu, and Yu-Gang Jiang, “Pixel2mesh: Generating 3D for 4d temporal medical image generation,” in Medical Im-
mesh models from single RGB images,” in Proceedings of age Computing and Computer Assisted Intervention–MICCAI
the European Conference on Computer Vision (ECCV), 2018, 2022: 25th International Conference, Singapore, September
pp. 52–67. 18–22, 2022, Proceedings, Part I. Springer, 2022, pp. 539–
548.
[6] Christopher B Choy, Danfei Xu, Junyoung Gwak, Kevin Chen,
and Silvio Savarese, “3D-R2N2: A unified approach for single [20] Ho, Jain, and Abbeel, “Denoising diffusion probabilistic mod-
and multi-view 3D object reconstruction,” in Proceedings of els,” Adv. Neural Inf. Process. Syst.
the European Conference on Computer Vision (ECCV), 2016, [21] Greta Simionato, Konrad Hinkelmann, Revaz Chachanidze,
pp. 628–644. Paola Bianchi, Elisa Fermo, Richard van Wijk, Marc Leonetti,
[7] Daniel Watson, William Chan, Ricardo Martin-Brualla, Christian Wagner, Lars Kaestner, and Stephan Quint, “Red
Jonathan Ho, Andrea Tagliasacchi, and Mohammad Norouzi, blood cell phenotyping from 3D confocal images using artifi-
“Novel view synthesis with diffusion models,” arXiv preprint cial neural networks,” PLOS Computational Biology, vol. 17,
arXiv:2210.04628, 2022. no. 5, pp. 1–17, 2021.
[22] Dominik Waibel, Niklas Kiermeyer, Scott Atwell, Bastian
[8] Haoqiang Fan, Hao Su, and Leonidas J. Guibas, “A point set
Rieck, Matthias Meier, and Carsten Marr, “Datasets for 3D
generation network for 3D object reconstruction from a single
shape reconstruction from 2D microscopy images,” Sept. 2022.
image,” in Proceedings of the IEEE Conference on Computer
Vision and Pattern Recognition (CVPR), 2017. [23] Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet,
Mohammad Norouzi, and Tim Salimans, “Cascaded diffusion
[9] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay models for high fidelity image generation.,” J. Mach. Learn.
Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Res., vol. 23, pp. 47–1, 2022.
Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes,
et al., “Photorealistic text-to-image diffusion models with deep [24] Alexandre Lacoste, Alexandra Luccioni, Victor Schmidt, and
language understanding,” arXiv preprint arXiv:2205.11487, Thomas Dandres, “Quantifying the carbon emissions of ma-
2022. chine learning,” arXiv preprint arXiv:1910.09700, 2019.