Home

Awesome

CounTR

Official PyTorch implementation for CounTR. Details can be found in the paper. [Paper] [Project page]

Thanks to @tamnguyenvan, now we can use CounTR in an easier way by using this library!

PWC

PWC

<img src=img/arch.png width="80%"/>

Contents

Preparation

1. Download datasets

In our project, the following datasets are used. Please visit following links to download datasets:

In fact, we use CARPK by importing hub package. Please click here for more information.

2. Download required python packages:

The following packages are suitable for NVIDIA GeForce RTX 3090.

pip install torch==1.10.0+cu111 torchvision==0.11.0+cu111 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install timm==0.3.2
pip install numpy
pip install matplotlib tqdm 
pip install tensorboard
pip install scipy
pip install imgaug
pip install opencv-python
pip3 install hub

CounTR Train

Please modify your work directory and dataset directory in the following train files.

Taskmodel filetrain file
Pretrain on FSC147models_mae_noct.pyFSC_pretrain.py
Finetune on FSC147models_mae_cross.pyFSC_finetune_cross.py
Finetune on CARPKmodels_mae_cross.pyFSC_finetune_CARPK.py

Pretrain on FSC147

CUDA_VISIBLE_DEVICES=0 python FSC_pretrain.py \
    --epochs 500 \
    --warmup_epochs 10 \
    --blr 1.5e-4 --weight_decay 0.05

Finetune on FSC147

CUDA_VISIBLE_DEVICES=0 nohup python -u FSC_finetune_cross.py \
    --epochs 1000 \
    --blr 2e-4 --weight_decay 0.05  >>./train.log 2>&1 &

Finetune on CARPK

CUDA_VISIBLE_DEVICES=0 nohup python -u FSC_finetune_CARPK.py \
    --epochs 1000 \
    --blr 2e-4 --weight_decay 0.05  >>./train.log 2>&1 &

CounTR Inference

Please modify your work directory and dataset directory in the following test files.

Taskmodel filetest file
Test on FSC147models_mae_cross.pyFSC_test_cross.py
Test on CARPKmodels_mae_cross.pyFSC_test_CARPK.py

Test on FSC147

CUDA_VISIBLE_DEVICES=0 nohup python -u FSC_test_cross.py >>./test.log 2>&1 &

Test on CARPK

CUDA_VISIBLE_DEVICES=0 nohup python -u FSC_test_CARPK.py >>./test.log 2>&1 &

Also, demo.py is a small demo used for testing on a single image.

CUDA_VISIBLE_DEVICES=0 python demo.py

Fine-tuned weights

benchmarkMAERMSElink
FSC14711.95 (Test set)91.23 (Test set)weights
CARPK5.757.45weights

Visualisation

<img src=img/goodpred.png width="75%"/>

Citation

@inproceedings{liu2022countr,
  author = {Chang, Liu and Yujie, Zhong and Andrew, Zisserman and Weidi, Xie},
  title = {CounTR: Transformer-based Generalised Visual Counting},
  booktitle={British Machine Vision Conference (BMVC)},
  year = {2022}
}

Acknowledgements

We borrowed the code from

Thanks @GioFic95 for adding the function of using external exemplars, more predictions images, more parametrized inference and so on.

If you have any questions about our code implementation, please contact us at liuchang666@sjtu.edu.cn