Spatial Diffusion for Cell Layout Generation

TL;DR

This work addresses the challenge of generating realistic cell layouts to augment pathology-image datasets for cell detection by introducing a spatial-pattern-guided diffusion framework. It represents cells as 3×3 markers and condition the diffusion process on both counting categories and spatial density maps, concatenating layout and density inputs via ${x}_0 = \text{concat}({x}_0^p, {x}_0^d)$; the forward process follows $q({x}_t|{x}_{t-1}) = \mathcal{N}({x}_t; \sqrt{1-\beta_t}{x}_{t-1}, \beta_t \boldsymbol{I})$ with a denoising reverse $p_\theta({x}_{t-1}|{x}_t) = \mathcal{N}({x}_{t-1}; \mu_\theta({x}_t, t), \sigma_t^2 \boldsymbol{I})$, and pathology images are generated conditioned on the layout via $p(I|{x}^p_g)$. The study compares KDE, GMM, and GMCM as spatial feature extractors, finding that GMMbest captures the BRCA-M2C layout clusters and yields the best Spatial-FID, and demonstrates that the generated layouts and images can significantly boost state-of-the-art cell-detection methods. Quantitatively, spatial-FID correlates with improved downstream detection performance, validating the practical utility of layout-guided image generation for augmenting pathology datasets. The authors release code and show that their approach yields realistic, density-aware layouts and high-quality pathology images suitable for detector training.

Abstract

Generative models, such as GANs and diffusion models, have been used to augment training sets and boost performances in different tasks. We focus on generative models for cell detection instead, i.e., locating and classifying cells in given pathology images. One important information that has been largely overlooked is the spatial patterns of the cells. In this paper, we propose a spatial-pattern-guided generative model for cell layout generation. Specifically, a novel diffusion model guided by spatial features and generates realistic cell layouts has been proposed. We explore different density models as spatial features for the diffusion model. In downstream tasks, we show that the generated cell layouts can be used to guide the generation of high-quality pathology images. Augmenting with these images can significantly boost the performance of SOTA cell detection methods. The code is available at https://github.com/superlc1995/Diffusion-cell.

Figures (4)

• Figure 1: Overview of proposed layout generation framework. In order to reduce the difficulty of learning spatial distributions behind object layouts, we propose a counting-category conditioned object layout generation framework. Meanwhile, the spatial density maps are incorporated into the training and generation process to help diffusion model learn and generate the spatial distribution of object layouts. For clearance, we only show one cell type's cell layout and spatial density map. ${\boldsymbol{x}}_t^p$ and ${\boldsymbol{x}}_t^d$ are the cell layout and spatial density map at time step $t$.
• Figure 2: Illustration of cell layout conditioned diffusion model for generating pathology images with ground truth.
• Figure 3: Qualitative results generated by our layout and image generation framework for cell detection. Rows 1st, 2nd, and 3rd are from the counting categories 0, 2, and 4, respectively. Tumor, lymphocyte, and stromal cells are marked by green, red, and blue marks.
• Figure 4: The three image pairs in the left top, right top, left medium, right medium, left bottom, and right bottom are from the training set, the generation of counting categories 0, 1, 2, 3, and 4, respectively. The images in 1st and 4th columns are pathology images. The cell layouts are in 2nd and 5th columns. The generated spatial density maps are in 3rd and 6th columns.