SAVE: Self-Attention on Visual Embedding for zero-shot generic object counting

Environment

Here is the code to setup the conda environment:

git clone https://github.com/AhmedZgaren/Save.git
cd Save
conda env create -f environment.yml
conda activate count

Download trained model

We provide pre-trained model weights in FSC147:

model

Store the pre-trained model in the pre-trained folder.

Demo

To test the model in an image use the code below:

python demo.py --i path\to\img

Train the model from scratch

Download FSC147 dataset

Use this link to download the FSC147 dataset and store images under the folder .\data\images\

Train your model

To train the model just run the following command:

python train.py

Test the trained model in different datasets

To test the model in the FSC147 validation set:

python valid.py -valid True

To test the model in the FSC147 test set:

python valid.py -test True

To test the model in the COCO validation set:

python valid.py -vcoco True

To test the model in the COCO test set:

python valid.py -tcoco True