cvlab-stonybrook / DM-Count

Code for NeurIPS 2020 paper: Distribution Matching for Crowd Counting.
MIT License
218 stars 52 forks source link

DM-Count

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

Arxiv | NeurIPS Processings

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

  1. Dataset download
  1. 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> 
  1. Training
python train.py --dataset <dataset name: qnrf, sha, shb or nwpu> --data-dir <path to dataset> --device <gpu device id>
  1. 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

A web interface to can be found here.

demo

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

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},
}