[Paper Review] A Diffusion Model Predicts 3D Shapes from 2D Microscopy Images

Dominik Waibel, Ernst Röell+3 more|arXiv (Cornell University)|Aug 30, 2022

Cell Image Analysis Techniques4 citations

TL;DR

This paper introduces DISPR, a diffusion model that predicts realistic 3D cell shapes from 2D microscopy images by conditioning on a single 2D fluorescent image. It significantly improves minority class classification in a red blood cell morphological feature task, increasing the macro F1 score from 55.2% to 72.2% through data augmentation using synthetic 3D reconstructions.

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_ ext{macro} = 55.2 \pm 4.6\%$ to $F1_ ext{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.

Motivation & Objective

• To address the challenge of slow, costly 3D confocal microscopy by enabling fast, realistic 3D shape prediction from 2D images.
• To overcome data scarcity and class imbalance in single-cell morphology classification by generating synthetic 3D shapes for data augmentation.
• To leverage the distributional nature of diffusion models to produce diverse, realistic 3D reconstructions from a single 2D input.
• To evaluate the utility of diffusion-based 3D shape prediction in improving downstream classification performance on imbalanced biological datasets.

Proposed method

• A denoising diffusion probabilistic model is trained to reverse a forward noising process that gradually adds Gaussian noise to 3D cell volumes over T steps.
• The model is conditioned on a 2D fluorescence microscopy image and a segmentation mask, which serve as priors during the denoising process.
• During inference, the model performs T denoising steps starting from pure noise, generating a 3D shape prediction at step t=0.
• The model is trained using a variational lower bound loss with a noise schedule that improves sampling efficiency.
• The architecture incorporates U-Net-based U-Net with skip connections and attention mechanisms to model long-range dependencies in 3D volumes.
• Synthetic 3D shapes are generated for data augmentation by predicting five 3D reconstructions per 2D image for minority classes.

Fig. 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 2 ). Note that the zoom level in the visualizations is not uniform.

Experimental results

Research questions

• RQ1 Can a diffusion model generate realistic 3D cell shapes from a single 2D microscopy image?
• RQ2 Does the distributional output of a diffusion model improve data augmentation for imbalanced single-cell classification tasks?
• RQ3 How do morphological features of DISPR-generated 3D shapes compare to ground truth in terms of realism and accuracy?
• RQ4 To what extent does using DISPR-generated data improve macro F1 scores in minority class classification?

Key findings

• The macro F1 score for classifying red blood cell morphologies improved from 55.2% ± 4.6% to 72.2% ± 4.9% when augmenting training data with features from DISPR-generated 3D shapes.
• The weighted F1 score increased from 85.7% ± 2.1% to 88.6% ± 1.1%, indicating improved overall classification performance.
• The number of correctly classified keratocytes increased from 14 to 24, and multilobate cells from 1 to 8 after data augmentation.
• DISPR-generated shapes exhibit lower relative surface roughness and curvature error compared to state-of-the-art deterministic models, indicating higher morphological fidelity.
• The model successfully captures class-specific morphological feature distributions from the ground truth dataset in its generated shape distribution.
• Despite higher inference time (4:45 minutes per prediction), DISPR outperforms deterministic autoencoders like SHAPR in generating diverse, realistic 3D reconstructions.

Fig. 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 a