This repository contains the code for the paper "Learning to Count without Annotations" published at CVPR'24 and allows training counting models in a self-supervised manner using Self-Collages. We provide an example notebook to experiment with the proposed method. Details can be found in the paper. [Paper]

While recent supervised methods for reference-based object counting continue to improve the performance on benchmark datasets, they have to rely on small datasets due to the cost associated with manually annotating dozens of objects in images. We propose UnCounTR, a model that can learn this task without requiring any manual annotations. To this end, we construct "Self-Collages", images with various pasted objects as training samples, that provide a rich learning signal covering arbitrary object types and counts. Our method builds on existing unsupervised representations and segmentation techniques to successfully demonstrate for the first time the ability of reference-based counting without manual supervision. Our experiments show that our method not only outperforms simple baselines and generic models such as FasterRCNN and DETR, but also matches the performance of supervised counting models in some domains.

The model is able to correctly predict the number of objects for clearly separated and slightly overlapping instances (see Subfigures a and c). The model also successfully identifies the object type of interest, e.g. in Subfigure b the density map correctly highlights the strawberries rather than blueberries. One limitation of our model is partial or occluded objects. For example, in Subfigure d the prediction missed a few burgers which are possibly the ones partially shown on the edge. However, partial or occluded objects are also challenging and ambiguous for humans.

The construction process of the Self-Collages is implemented in SelfCollageDataset.py. The UnCounTR model is implemented in UnCounTRModel.py.

To reproduce the results of the paper, the following steps are required:

1. Cloning this repository
2. Downloading datasets
3. Downloading pretrained weights
4. Cloning third-party code
5. Installing environments

FSC147_384_V2
└── annotation_FSC147_384.json
└── ImageClasses_FSC147.txt
└── Train_Test_Val_FSC_147.json
└── gt_density_map_adaptive_384_VarV2
│   └── *.npy
│    ...
└── images_384_VarV2
    └── *.jpg
     ...

MSO
└── imgIdx.mat
└── img
    └── *.jpg
     ...

CARPK
└── Annotations
│   └── *.txt
│    ...
└── Images
│   └── *.png
│    ...
└── ImageSets
    └── test.txt
    └── train.txt

SUN397
└── Partitions
│   └── Testing_0*.txt
│    ...
└── SUN397
    └── a
        └── abbey
            └── *.jpg
             ...
         ...
     ...

ImageNet
└── ILSVRC2012_devkit_t12
│    ...
└── train
│   └── n*
│       └── *.JPEG
│        ...
│    ...
└── val
    └── n*
        └── *.JPEG
         ...
     ...

stylegan-oriented
└── 00000
    └── *.jpg
         ...
     ...

To obtain object segmentations, we use selfmask. Download the pretrained weights for the model with 20 queries (file name: selfmask_nq20.pt) and place the file in SelfCollages/data/.

If you want to run the UnCounTR model with a pretrained Leopart backbone, download the model weights from here and place them in SelfCollages/data/. Specifically, the filenames should be leopart_vitb8.ckpt and leopart_vits16.ckpt.

Create a directory called SelfCollages/src/third_party and clone the following repositories into it:

• CounTR
• selfmask

To run this code, you need to install the Anaconda environment specified in env.yml (for CPU only) or env_gpu.yml (for GPU). This can be done using one of the following commands when in the root directory of this repository:

conda env create -f env_files/env.yml or conda env create -f env_files/env_gpu.yml

With the object mask generation being the only exception, all commands require the previously installed conda environment. It can be activated using:

conda activate uncountr

To generate object masks, we use selfmask. Make sure to install the necessary requirements for this code as mentioned in the corresponding repository.

Instead of training UnCounTR from scratch, you can download the pretrained model here. Extract the zip-file and place the model folder in SelfCollages/runs/. The folder should contain the following files:

uncountr_model
└── args.pt
└── UnCounTRModel.pt

All commands should be executed from the root directory of this repository. To get more information about optional arguments, use the --help flag.

To train UnCounTR from scratch, the object masks of the unlabelled ImageNet images have to be generated first followed by the training step itself.

Generate object masks This step creates object masks for the images in the ImageNet dataset which are saved in /path/to/SelfCollages/data/ImageNet/segmentations/selfmask. Unlike all other commands, this step must be executed in the selfmask environment.

python mask_generator.py --data_dir=/path/to/SelfCollages/data --img_net_path=/path/to/SelfCollages/data/ImageNet

Train UnCounTR model This step trains the UnCounTR model using Self-Collages. The trained model will be saved in /path/to/SelfCollages/runs/model_name.

python train_UnCounTR.py

Evaluate UnCounTR The UnCounTR model can be tested on different datasets and subsets using:

python evaluate_UnCounTR.py  --model_dir=/path/to/SelfCollages/runs/model_name --data_path=/path/to/SelfCollages/data/FSC147_384_V2 --weights_dir=/path/to/SelfCollages/data --output_dir=/path/to/output_directory --dataset_type=dataset_type

Aggregate results After evaluation, the following command allows to compute additional metrics and to calculate the performance on different subsets:

python aggregate_evaluation_results.py --eval_results_path=/path/to/output_directory/dataset_type/subdir --dataset_type=dataset_type

Evaluate baselines To compare UnCounTR to several baselines, they can be evaluated using:

python test_baselines.py --img_size=384 --batch_size=32

After training UnCounTR from scratch or downloading the pretrained model, it can be used for self-supervised semantic counting.

python self_supervised_semantic_counting.py --model_dir=/path/to/model --img_dir=/path/to/imgs

The example notebook contains the commands described above to train and evaluate UnCounTR as well as the necessary steps to experiment with self-supervised semantic counting. To start the notebook, use the following command:

jupyter notebook SelfCollages.ipynb

You might want to increase the maximum amount of memory that can be used by the notebook. This can be done with the argument --NotebookApp.max_buffer_size=X where X is the maximum amount of memory in bytes.

If you find this repository useful, please consider citing our paper:

@inproceedings{knobel2024learning,
      title={Learning to Count without Annotations}, 
      author={Lukas Knobel and Tengda Han and Yuki M. Asano},
      booktitle={CVPR},
      year={2024}
}

The code is licensed under the MIT License except for code taken from other sources, where we specify the source at the beginning of the file. For the pretrained weights, please refer to the license specified in the DINO repository.