This repository holds the official Pytorch implementation for the paper A Novel Unified Architecture for Low-Shot Counting by Detection and Segmentation accepted at NeurIPS 2024.

Low-shot object counters estimate the number of objects in an image using few or no annotated exemplars. Objects are localized by matching them to prototypes, which are constructed by unsupervised image-wide object appearance aggregation. Due to potentially diverse object appearances, the existing approaches often lead to overgeneralization and false positive detections. Furthermore, the best-performing methods train object localization by a surrogate loss, that predicts a unit Gaussian at each object center. This loss is sensitive to annotation error, hyperparameters and does not directly optimize the detection task, leading to suboptimal counts.We introduce GeCo, a novel low-shot counter that achieves accurate object detection, segmentation, and count estimation in a unified architecture. GeCo robustly generalizes the prototypes across objects appearances through a novel dense object query formulation. In addition, a novel counting loss is proposed, that directly optimizes the detection task and avoids the issues of the standard surrogate loss. GeCo surpasses the leading few-shot detection-based counters by ~25% in the total count MAE, achieves superior detection accuracy and sets a new solid state-of-the-art result across all low-shot counting setups.

To install the required dependencies, run the following command:

conda create -n geco_test python=3.8
conda activate geco_test
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install matplotlib

To run the demo, you need to download the pretrained weights and put them in the MODEL_folder.

Run the demo:

python demo.py --image_path ./material/4.jpg --output_masks

To evaluate GeCo on FSC147, install also:

pip install tqdm
pip install pycocotools
pip install scipy
python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'

download all the required data:

1. The original FSC147 dataset from Learning to Count Everything (put in the DATA_folder),

2. Box annotations for validation and test split from Counting-DETR (put in the DATA_folder/annotations),

3. Pretrained weights (put in the MODEL_folder).

and compute density maps:

python utils/data.py --data_path DATA_folder

(Need to compute density maps due to FSCD147 incompatibility with the original FSC147 annotations.)

Run inference on FSC147:

python evaluate.py --data_path DATA_folder --model_path MODEL_folder

Run bbox evaluation on FSC147:

python evaluate_bboxes.py --data_path DATA_folder

To train the model, follow the steps for evaluation on FSC147, correct paths in train.sh and pretrain.sh, download box annotations for train split, put them in DATA_folder, SAM-HQ pretrained weights, put them in MODEL_folder and run the following commands:

First, generate density maps:

python utils/data.py

First run pretraining:

sbatch pretrain.sh

then run the main training:

sbatch train.sh

@article{pelhan2024novel,
  title={A Novel Unified Architecture for Low-Shot Counting by Detection and Segmentation},
  author={Pelhan, Jer and Lukezic, Alan and Zavrtanik, Vitjan and Kristan, Matej},
  journal={Advances in Neural Information Processing Systems},
  volume={37},
  pages={66260--66282},
  year={2024}
}