CoDi - Counting with Diffusion

CoDi - Exemplar-Conditioned Diffusion Model for Low-Shot Counting

Grega Šuštar, Jer Pelhan, Alan Lukežič, Matej Kristan

University of Ljubljana, Faculty of Computer and Information Science

Overveiw | Get Started | Reference | Questions

🚧 UNDER CONSTRUCTION 🚧

Overview

This repository contains the official PyTorch implementation of CoDi - Exemplar-Conditioned Diffusion Model for Low-Shot Counting

CoDi generates high-quality density maps using latent diffusion model, enabling accurate object localization via non-maxima suppression.

The key innovation is an exemplar-based conditioning module that extracts and adapts object prototypes within the denoising network’s intermediate layers, significantly improving location estimation.

generation.mp4

Get Started

Environment setup

Create a Python environment and install required packages:

conda create -n codi python=3.10
conda activate codi

pip install -r requirements.txt

Download the datasets

Download the FSC147 and MCAC datasets as instructed in their official repositories. For the FSC147 datasets, make sure to also download the annotation_FSC147_384.json and Train_Test_Val_FSC_147.json and place them alongside the image directory (images_384_VarV2) in the directory of your choice.

Checkpoints:

The checkpoints are available on Google drive.

Training, Sampling and Evaluation

Before running any of the following scripts, make sure to replace the datadir path in the config files to point to your data directory folder. For additional arguments check the train.py, sample.py and eval.py scripts.

Training

./run_train.sh configs/config_codi.yaml

Sampling

For sampling on the MCAC dataset, create a new folder sample (in the same directory as train, val and test) with all the images you would like to sample on (or just copy the contents of val and test into it).

python sample.py \
    --expdir ../CoDi_ckpts/codi/ \
    --checkpoint model000150_codi.pt \
    --use_fp16 \
    --allow_resizing \
    --allow_ttn \
    --ids image-ids \                   # 7 6653 1123 ... 
    --timestep_respacing 25

Evaluation

python eval.py 
    --expdir ../CoDi_ckpts/codi/ \
    --checkpoint model000150_codi.pt \
    --use_fp16 \
    --allow_resizing \
    --allow_ttn \
    --timestep_respacing 25

Results

Reference

If you found this work useful, consider citing our paper and giving this repo a ⭐ 😃

@misc{sustar2025codi,
  author={Grega Šuštar, Jer Pelhan, Alan Lukežič, Matej Kristan},
  title={CoDi - Exemplar-Conditioned Diffusion Model for Low-Shot Counting}, 
  year={2025},
  eprint={2512.20153},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2512.20153}, 
}

Questions

For issues or questions, please open a GitHub issue or email the author directly.