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.
Python 3.x
Pytorch >= 1.2
For other libraries, check requirements.txt.
- 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 on UCF-QNRF, NWPU, Shanghaitech part A and B can be found in pretrained_models folder or Google Drive
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},
}