[论文解读] 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 Techniques被引用 4

一句话总结

本文提出DISPR，一种扩散模型，通过仅依赖单张2D荧光图像作为条件，从2D显微镜图像预测逼真的3D细胞形态。该方法通过使用合成3D重建进行数据增强，在红细胞形态特征分类任务中显著提升了少数类分类性能，将宏平均F1分数从55.2%提升至72.2%。

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.

研究动机与目标

• 为解决3D共聚焦显微镜成像速度慢、成本高的问题，实现从2D图像快速生成逼真3D形态预测。
• 通过生成合成3D形态用于数据增强，克服单细胞形态分类中数据稀缺与类别不平衡的问题。
• 利用扩散模型的分布特性，从单张2D输入生成多样化且逼真的3D重建。
• 评估基于扩散模型的3D形态预测在提升不平衡生物数据集下游分类性能方面的实用性。

提出的方法

• 训练一个去噪扩散概率模型，以逆转在T步内逐步向3D细胞体积添加高斯噪声的前向加噪过程。
• 模型以2D荧光显微镜图像和分割掩码作为条件，二者在去噪过程中充当先验信息。
• 推理阶段，模型从纯噪声开始执行T步去噪，最终在t=0步生成3D形态预测。
• 采用变分下界损失进行训练，并使用优化的噪声调度以提升采样效率。
• 网络架构采用带有跳跃连接和注意力机制的U-Net结构，以建模3D体积中的长程依赖关系。
• 通过为少数类的每张2D图像预测五个3D重建，生成合成3D形态用于数据增强。

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.

实验结果

研究问题

• RQ1 扩散模型能否从单张2D显微镜图像生成逼真的3D细胞形态？
• RQ2 扩散模型的分布输出是否能提升少数类分类任务中数据增强的效果？
• RQ3 DISPR生成的3D形态在形态学特征上与真实数据相比，其逼真度和准确性如何？
• RQ4 使用DISPR生成的数据在多大程度上提升了少数类分类的宏平均F1分数？

主要发现

• 在训练数据中加入DISPR生成的3D形态特征后，红细胞形态分类的宏平均F1分数从55.2% ± 4.6%提升至72.2% ± 4.9%。
• 加权F1分数从85.7% ± 2.1%提升至88.6% ± 1.1%，表明整体分类性能得到改善。
• 数据增强后，正确分类的成纤维样细胞（keratocytes）数量从14例增至24例，多叶状细胞（multilobate cells）从1例增至8例。
• 与当前最先进的确定性模型相比，DISPR生成的形态具有更低的相对表面粗糙度和曲率误差，表明其形态保真度更高。
• 模型成功在生成的形态分布中复现了真实数据集中各类细胞特有的形态学特征分布。
• 尽管推理时间较长（每预测耗时4:45分钟），DISPR在生成多样化、逼真3D重建方面仍优于如SHAPR等确定性自编码器。

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