SAFECount: Few-shot Object Counting with Similarity-Aware Feature Enhancement

repository·main·Indexed 40 days ago

https://github.com/zhiyuanyou/safecount

A PyTorch implementation of the WACV 2023 paper 'Few-shot Object Counting with Similarity-Aware Feature Enhancement'. It provides tools for few-shot object counting across original FSC147, cross-validation, cross-dataset generalization (e.g., FSC147 to CARPK), and class-specific counting using datasets like PUCPR+, UCSD, Mall, and ShanghaiTech.

Tokens
1.3K
Snippets
3
Records
6
Agent score
0%

What's inside SAFECount

  1. How datasets and models cooperate in SAFECount

    main

    The choice of dataset class and model depends on your specific counting scenario:

    ScenarioDataset ClassModel ClassDescription
    FSC147 Stylecustom_dataset.pysafecount.pySupport images are parts of the query image, annotated by bounding boxes.
    Class-specificcustom_exemplar_dataset.pysafecount_exemplar.pySupport images are sampled and fixed, not part of the query image (e.g., CARPK, Mall).
    Cross-datasetcustom_dataset.pysafecount_crossdataset.pyPretraining on FSC147 and finetuning on class-specific datasets.
  2. Quick Start: Cross-dataset Generalization (FSC147 to CARPK)

    main

    To perform cross-dataset generalization (pretraining on FSC147 and finetuning on CARPK):

    1. Prepare datasets: Set up both FSC147 and CARPK dataset directories following the structures required for their respective settings.
    2. Navigate to the experiment directory: cd experiments/FSC147_to_CARPK/.
    3. Run workflows: Execute pretraining, finetuning, or evaluation using the standard training/eval scripts.
  3. Quick Start: FSC147 Cross Validation Setting

    main

    To run experiments using cross-validation (e.g., 3-shot with fold0 as the validation set):

    1. Prepare the FSC147 dataset as described in the Original Setting guide.
    2. Navigate to the specific fold directory: cd experiments/FSC147_fold/3shot/fold0/.
    3. Execute training or evaluation using Slurm or torch.distributed.launch commands.

    Note: Ensure config.saver.load_path is set to load checkpoints during eval/test.

    # For Slurm group
    sh train.sh #NUM_GPUS #PARTITION
    sh eval.sh #NUM_GPUS #PARTITION
    
    # For torch.distributed.launch
    sh train_torch.sh #NUM_GPUS #GPU_IDS
    sh eval_torch.sh #NUM_GPUS #GPU_IDS
  4. Quick Start: FSC147 Original Setting

    main

    To run the SAFECount model in the original FSC147 setting, follow these steps:

    1. Prepare the FSC147 dataset: Download the FSC147 dataset and organize it into ./data/FSC147_384_V2/ with the following structure:

      • images_384_VarV2
      • gt_density_map_adaptive_384_VarV2
      • train.json, val.json, test.json, fold0.json, fold1.json, fold2.json, fold3.json
    2. Navigate to the experiment directory: cd ./experiments/FSC147/.

    3. Execute training, evaluation, or testing using either Slurm or torch.distributed.launch.

    Note: When performing evaluation or testing, you must set config.saver.load_path in your configuration to point to your checkpoints.

    # For Slurm group
    sh train.sh #NUM_GPUS #PARTITION
    sh eval.sh #NUM_GPUS #PARTITION
    sh test.sh #NUM_GPUS #PARTITION
    
    # For torch.distributed.launch (e.g., 4 GPUs on IDs 1,3,4,6)
    sh train_torch.sh 4 1,3,4,6
    sh eval_torch.sh 4 1,3,4,6
    sh test_torch.sh 4 1,3,4,6
  5. Prepare Datasets for Class-specific Counting

    main

    Class-specific counting requires specific directory structures and ground-truth density map generation. For most datasets (CARPK, PUCPR+, UCSD, Mall, ShanghaiTech), you must run the provided gen_gt_density.py script after downloading and unzipping the data.

    CARPK & PUCPR+

    Structure:

    • data/CARPK_devkit/ (or PUCPR+_devkit/)
      • Images
      • gen_gt_density.py
      • train.json, test.json, exemplar.json

    Action: Run python gen_gt_density.py to generate maps in ./data/[Dataset]_devkit/gt_density_map/.

    UCSD

    Structure:

    • data/UCSD/ucsdpeds_vidf/ (Merge images from the 10 vidf1_33_00x.y directories into this folder).
    • data/UCSD/gen_gt_density.py
    • data/UCSD/train.json, test.json, exemplar.json, mask.png

    Action: Run python gen_gt_density.py to generate maps in ./data/UCSD/gt_density_map/.

    Mall

    Structure:

    • data/Mall/
      • frames
      • gen_gt_density.py
      • train.json, test.json, exemplar.json, mask.png

    Action: Run python gen_gt_density.py to generate maps in ./data/Mall/gt_density_map/.

    ShanghaiTech

    Structure:

    • data/ShanghaiTech/part_A/ (and part_B/)
      • train_data/images
      • test_data/images
      • gen_gt_density.py
      • train.json, test.json, exemplar.json

    Action: Run python gen_gt_density.py twice (once for part_A and once for part_B) to generate maps for both train_data and test_data subdirectories.

    # Example for CARPK/PUCPR+
    python gen_gt_density.py
  6. Troubleshoot CUDA Out of Memory errors

    main

    If you encounter CUDA Out of Memory errors, try the following adjustments in your configuration:

    1. Reduce Image Size: Decrease config.dataset.input_size.
    2. Reduce Exemplars: Decrease config.dataset.shot (for FSC147) or config.dataset.exemplar.num_exemplar (for class-specific counting).
    3. Increase Stride: Set a larger config.net.kwargs.backbone.out_stride.
      • Note: If you increase the stride, you must also modify the Regressor in models.utils.py to upsample the feature back to the original image size.