VAE-Diff: Mask-Conditioned Diffusion for Synthetic Cytological Image Generation

Reproducibility repository for the paper "Mask-Conditioned Diffusion for Synthetic Cytological Image Generation", submitted to Biomedical Signal Processing and Control.

This repository consolidates the full pipeline: a VAE that generates single-cell masks (VAE-Diff), the mask-conditioned diffusion renderer, the unconditional baselines (Diff-I, Diff-IM), the conditional GAN baseline (cGAN), and the evaluation and analysis code that produces every table and figure in the paper.

Data availability. The data cannot be shared publicly due to ethical restrictions (Humanitas Research Hospital, approval ICH2782) but may be available upon reasonable request subject to institutional approval. This repository releases code only — it contains no patient data, image datasets, or trained model weights. The two figures under docs/ are the paper's illustrative examples.

Pipeline

Two-stage approach for synthetic cytology image generation. Step 1: the VAE is trained on individual cell masks, cropped from the segmentation masks of real patches, and learns to reconstruct a single cell from a 256-dimensional latent code; sampling that latent yields new single-cell masks. Step 2: the diffusion model is trained on real patches, with the binary mask concatenated to the RGB image as a fourth input channel, so it learns to denoise an image conditioned on where the cells are. Generation: single-cell masks are sampled from the VAE and assembled into a full 256x256 patch by the mask-composition step, which draws the number of cells from the empirical distribution measured on real patches, places each mask at a random non-overlapping position, and allows border masks to be cropped. The composed mask then conditions the diffusion model, which renders the image; the pipeline therefore emits a synthetic image and its exact segmentation mask together, with no annotation step.

Example generations

Mask-conditioned generation, compared under a fixed mask. Each row is a single VAE-generated mask (first column) rendered independently by VAE-Diff (second column) and by the conditional GAN baseline (third column); the mask is identical between the two, so the generator is the only variable. The fourth column, past the dashed rule, is an independent real patch matched only on the number of masked cells. Rows are ordered by mask occupancy, from crowded fields with overlapping cells to sparse ones. Masks cover nucleated cells only; the unmasked erythrocytes form the background field. All patches are shown at their native 256x256 resolution.

Downstream utility: segmentation

A segmentation model trained only on synthetic data works on real data. A U-Net trained purely on VAE-Diff synthetic image/mask pairs reaches a median Dice of 0.941 on held-out real test patches, against 0.945 for a U-Net trained on real data. Each row is one real test patch; contours are ground truth (cyan), the real-trained U-Net (grey), and the synthetic-only U-Net (green), with per-prediction Dice. Rows are the patches whose synthetic-trained Dice lies closest to the test-set median, so they are representative rather than best-case.

Layout

vae-diff-cytology/
├── pipeline/          orchestration, generation, evaluation and figure/table code (this work)
├── external/          the four upstream engine repos, as pinned submodules (see below)
│   ├── MaskConditionedDiffusion/    diffusion train + generation (cond / uncond / 4-channel)
│   ├── VAE4Mask/                    VAE for single-cell masks + composition/placement
│   ├── SyntheticCitologyValidation/ FID-infinity, KID, feature-Wasserstein metrics
│   └── CellSegmentation/            U-Net segmentation training + Dice (TSTR)
├── docs/              figures used in this README
├── CITATION.cff
├── LICENSE            applies to pipeline/ only (see "Licensing")
├── REPRODUCE.md       environment, data layout, entry points
├── localize.sh        rewrite pipeline/ placeholder paths to your layout
└── setup_submodules.sh

Engine repositories (submodules)

The four engines under external/ are included as submodules pinned to the exact commits used for the paper. They are forks of their upstream repositories; the fork lineage on GitHub preserves the original authorship and history. Run setup_submodules.sh after cloning, or git submodule update --init if already wired.

Reproducing the paper

See REPRODUCE.md for the environment, directory layout, external dependencies, and entry points. In short: git submodule update --init, create the cytogen conda environment, run ./localize.sh to point the scripts at your paths, obtain the data (institutional approval), then use the entry points in pipeline/.

Licensing

• The code in pipeline/ is released under the MIT License (see LICENSE), this work.
• The external/ submodules are separate upstream repositories, included by reference (fork + submodule) and not relicensed here. They carry no upstream license; they are used under GitHub's Terms of Service, which permit forking of public repositories. The MIT License in this repository does not extend to them. Preserve upstream authorship (carried by the fork lineage and commit history) when reusing.