DM-Count

Official Pytorch implementation of the paper Distribution Matching for Crowd Counting (NeurIPS, spotlight).

We propose to use Distribution Matching for crowd COUNTing (DM-Count). In DM-Count, we use Optimal Transport (OT) to measure the similarity between the normalized predicted density map and the normalized ground truth density map. To stabilize OT computation, we include a Total Variation loss in our model. We show that the generalization error bound of DM-Count is tighter than that of the Gaussian smoothed methods. Empirically, our method outperforms the state-of-the-art methods by a large margin on four challenging crowd counting datasets: UCF-QNRF, NWPU, ShanghaiTech, and UCF-CC50.

Prerequisites

Python 3.x

Pytorch >= 1.2

For other libraries, check requirements.txt.

Getting Started

Dataset download

QNRF can be downloaded here
NWPU can be downloaded here
Shanghai Tech Part A and Part B can be downloaded here

Data preprocess

Due to large sizes of images in QNRF and NWPU datasets, we preprocess these two datasets.

python preprocess_dataset.py --dataset <dataset name: qnrf or nwpu> --input-dataset-path <original data directory> --output-dataset-path <processed data directory>

Training

python train.py --dataset <dataset name: qnrf, sha, shb or nwpu> --data-dir <path to dataset> --device <gpu device id>

Test

python test.py --model-path <path of the model to be evaluated> --data-path <directory for the dataset> --dataset <dataset name: qnrf, sha, shb or nwpu>

Pretrained models

Pretrained models on UCF-QNRF, NWPU, Shanghaitech part A and B can be found Google Drive. You could download them and put them in in pretrained_models folder.

Other resources

Web Demo

A web interface to can be found here.

Feel free to upload a image and try out the demo on a web browser. It is developed by Ali Abdalla from Gradio. Gradio is an open source library, which helps to create interfaces to make models more accessible. Thanks Ali and Gradio!

To launch a Gradio interface, run

 python demo.py

Kaggle Notebook

A Kaggle Notebook is developed by Selman Ozleyen. Thanks Selman!

References

If you find this work or code useful, please cite:

@inproceedings{wang2020DMCount,
  title={Distribution Matching for Crowd Counting},
  author={Boyu Wang and Huidong Liu and Dimitris Samaras and Minh Hoai},
  booktitle={Advances in Neural Information Processing Systems},
  year={2020},
}