A Diffusion Model Predicts 3D Shapes from 2D Microscopy Images

TL;DR

This work presents DISPR, a diffusion-based method for reconstructing 3D cell shapes from 2D microscopy by conditioning the diffusion process on the 2D input. It demonstrates superior geometric fidelity compared with baselines and prior 3D reconstruction methods, and shows DISPR can meaningfully augment training data for imbalanced red blood cell classes, boosting macro F1 scores. The approach yields multiple plausible 3D reconstructions per 2D image and highlights the utility of distribution-aware generative methods for inverse biomedical problems, albeit with higher inference time. Overall, DISPR offers a data-augmentation-enabled pathway to high-resolution 3D cellular morphology from limited 2D data, with potential applicability beyond red blood cells.

Abstract

Diffusion models are a special type of generative model, capable of synthesising new data from a learnt distribution. We introduce DISPR, a diffusion-based model for solving the inverse problem of three-dimensional (3D) cell shape prediction from two-dimensional (2D) single cell microscopy images. Using the 2D microscopy image as a prior, DISPR is conditioned to predict realistic 3D shape reconstructions. To showcase the applicability of DISPR as a data augmentation 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 predictions to the three minority classes improved the macro F1 score from $F1_\text{macro} = 55.2 \pm 4.6\%$ to $F1_\text{macro} = 72.2 \pm 4.9\%$. We thus demonstrate that diffusion models can be successfully applied to inverse biomedical problems, and that they learn to reconstruct 3D shapes with realistic morphological features from 2D microscopy images.

Figures (5)

• Figure 1: Diffusion steps of one forward pass through DISPR visualized during inference between step $t=999$ with the largest amount of Gaussian noise and the DISPR prediction of $t=0$, for a stromatocyte cell (see section \ref{['sec:Methods']}). Note that the zoom level in the visualizations is not uniform.
• Figure 2: DISPR is trained to denoise the 3D volume $\mathbf{x}_{b,t}$ containing added stochastic Gaussian noise, to obtain $\mathbf{x}_{b,t-1}$, thus reversing the noising step $q$. In each forward pass $p_{\theta}$, we constrain our 3D model with one 2D image $b$ containing a fluorecent image and a mask. During inference, the forward pass $p_{\theta}$ is repeated $T$ times to obtain the prediction $\mathbf{x}_{b,0}$ (see section \ref{['sec:Methods']}).
• Figure 3: Red blood cells exhibit similar morphological features between groundtruth (3rd column, yellow background) and DISPR predictions (white background). A 2D microscopy image (with a fluorescence and segmented channel, dark background) is used as the model's input to predict the 3D shapes.
• Figure 4: With the predictions of DISPR we obtain a lower relative volume error as compared to the extrapolations of two naive models, the cylinder and ellipsoid fit, as well as predictions of SHAPR Waibel2021-or and the predictions of SHAPR using the topological loss Waibel2022-ng. 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 times as many datapoints than the others.
• Figure 5: Using DISPR to oversample training data improves a feature-based random forest classification of single red blood cells. (a) For the training dataset, containing 128 manually extracted morphological features of the 3D groundtruth, we enhance the three smallest classes (multilobate, keratocyte, knizocyte) with features extracted from the DISPR's predictions, reducing class imbalance. (b) We improve the macro F1 score to $F1_{\text{macro}} = 72.2\pm4.9\%$, from $F1_{\text{macro}}$ = $55.2\pm4.6\%$, with the largest improvements obtained for minority classes (c).