Official PyTorch implementation for CounTR. Details can be found in the paper. [Paper] [Project page]
- Preparation
- CounTR train
- CounTR inference
- Fine-tuned weights
- Visualisation
- Citation
- Acknowledgements
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.
The following packages are suitable for NVIDIA GeForce RTX 3090. Works on Linux only：
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.4.5
pip install numpy==1.23.5
pip install matplotlib tqdm
pip install tensorboard
pip install scipy
pip install imgaug
pip install opencv-python
pip install pandas
pip3 install hub
My Version for Windows：
pip install torch==1.10.0+cu113 torchvision==0.11.0+cu113 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install timm==0.4.5
pip install numpy==1.23.5
pip install matplotlib tqdm
pip install tensorboard
pip install scipy
pip install imgaug
pip install opencv-python
pip install pandas
pip3 install hub
Please modify your work directory and dataset directory in the following train files.
| Task | model file | train file |
|---|---|---|
| Pretrain on FSC147 | models_mae_noct.py | FSC_pretrain.py |
| Finetune on FSC147 | models_mae_cross.py | FSC_finetune_cross.py |
| Finetune on CARPK | models_mae_cross.py | FSC_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 &
Please modify your work directory and dataset directory in the following test files.
| Task | model file | test file |
|---|---|---|
| Test on FSC147 | models_mae_cross.py | FSC_test_cross.py |
| Test on CARPK | models_mae_cross.py | FSC_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
| benchmark | MAE | RMSE | link |
|---|---|---|---|
| FSC147 | 11.95 (Test set) | 91.23 (Test set) | weights |
| CARPK | 5.75 | 7.45 | weights |
@article{liu2022countr,
author = {Chang, Liu and Yujie, Zhong and Andrew, Zisserman and Weidi, Xie},
title = {CounTR: Transformer-based Generalised Visual Counting},
journal = {arXiv:2208.13721},
year = {2022}
}
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 [email protected]
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent