• Code
  • Dataset
  • Demo
  • Model

DiffuMT

Mask-conditioned diffusion model for IRM microtubule images

A diffusion model that turns a segmentation mask into a realistic IRM microtubule image, for generating labeled training data.

DiffuMT generates realistic 256 × 256 px interference reflection microscopy (IRM) images of microtubules from a binary segmentation mask. Because the mask is known, every generated image comes with its label, which makes the model a source of labeled training data for microtubule segmentation.

It is a mask-conditioned denoising diffusion model (a 2D U-Net with about 113 million parameters) fine-tuned on IRM microtubule images. The released checkpoint was chosen by a diagnostic that balances realism, diversity and colour match, before the model starts to collapse to fewer modes later in training. The diagnostic is described in our paper, accepted as a spotlight at NLDL 2027. In a two-alternative forced-choice study, domain experts could not reliably tell the generated images from real recordings, and a segmentation model trained only on generated images performed as well as one trained on real data. The generated images are released as the DiffuMT dataset.

The preview shows real IRM images in the top row and, below each, the image DiffuMT generated from its mask. On the project page, you can draw a mask in the browser and watch the model sample an image step by step.

from diffusers import UNet2DModel, DDIMScheduler

unet = UNet2DModel.from_pretrained("HTW-KI-Werkstatt/DiffuMT", subfolder="unet")
scheduler = DDIMScheduler.from_pretrained("HTW-KI-Werkstatt/DiffuMT", subfolder="scheduler")

Note that the model expects masks scaled to {0, 1/255}, not {0, 1}; the model card explains how.