SAFECount

https://github.com/zhiyuanyou/safecount
Admin
SAFECount is a deep learning framework for few-shot object counting that uses similarity-aware...
Tokens:7,132
Snippets:62
Trust Score:9.3
Update:6 months ago
Show doc for...
Tokens:
Raw
### Configuration YAML for Model and Training

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Example structure of a YAML configuration file for defining hyperparameters related to dataset, model architecture, training, and evaluation.

```yaml
...
```

--------------------------------

### Train, Eval, or Test FSC147 (torch.distributed.launch)

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

These commands are for training, evaluating, or testing on the FSC147 dataset using torch.distributed.launch. Specify the total number of GPUs and their IDs. For example, to train on GPUs 1, 3, 4, and 6 (4 GPUs total), use: sh train_torch.sh 4 1,3,4,6. Remember to set config.saver.load_path for evaluation or testing.

```bash
sh train_torch.sh #NUM_GPUS #GPU_IDS
```

```bash
sh eval_torch.sh #NUM_GPUS #GPU_IDS
```

```bash
sh test_torch.sh #NUM_GPUS #GPU_IDS
```

--------------------------------

### Load CustomDataset for FSC147 Data

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Loads images with in-image exemplar boxes for few-shot counting. Each sample contains the query image, ground truth density map, and bounding boxes marking example objects. Ensure data directories and meta files are correctly specified.

```python
import torch
import torchvision.transforms as transforms
from datasets.custom_dataset import CustomDataset
from datasets.base_dataset import BaseTransform

# Setup transforms
input_size = [512, 512]
transform_fn = BaseTransform(
    input_size,
    hflip={"prob": 0.5},      # Random horizontal flip
    vflip=False,
    rotate=False,
    gamma={"range": [0.8, 1.25], "prob": 0.5},
    gray=False,
)
normalize_fn = transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225],
)

# Create dataset
dataset = CustomDataset(
    img_dir="./data/FSC147_384_V2/images_384_VarV2/",
    density_dir="./data/FSC147_384_V2/gt_density_map_adaptive_384_VarV2/",
    meta_file="./data/FSC147_384_V2/train.json",
    shot=3,                    # Number of exemplar boxes to use (None = use all)
    transform_fn=transform_fn,
    normalize_fn=normalize_fn,
    colorjitter_fn=None,
)

# Access a sample
sample = dataset[0]
print(f"Image shape: {sample['image'].shape}")      # [C, H, W]
print(f"Density shape: {sample['density'].shape}")  # [1, H, W]
print(f"Boxes shape: {sample['boxes'].shape}")      # [M, 4] - [y1, x1, y2, x2]
print(f"Ground truth count: {sample['density'].sum():.1f}")
```

--------------------------------

### Load Pretrained Model Checkpoints

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Builds the network architecture and restores model weights from a checkpoint file.

```python
import torch
from models.safecount import build_network
from utils.misc_helper import load_state

# Build model
model = build_network(
    block=4,
    backbone={"type": "resnet18", "out_layers": [1, 2, 3], "out_stride": 4},
    pool={"type": "max", "size": [3, 3]},
    embed_dim=256,
    mid_dim=1024,
    head=8,
    dropout=0,
    activation="leaky_relu",
    exemplar_scales=[0.8, 1, 1.2],
)
model.cuda()

# Load checkpoint
checkpoint_path = "./checkpoints/ckpt_best.pth.tar"
best_mae, last_epoch = load_state(
    checkpoint_path,
    model,
    optimizer=None,  # Pass optimizer to restore training state
)
print(f"Loaded checkpoint from epoch {last_epoch}, best MAE: {best_mae:.2f}")

# Set to evaluation mode
model.eval()
```

--------------------------------

### Visualize Density Maps with Visualizer

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

The `build_visualizer` function creates a visualizer object for generating heatmap overlays. Use `vis_batch` to visualize predicted density maps on input images. Ensure `vis_dir` and `img_dir` are correctly specified.

```python
from utils.vis_helper import build_visualizer
import torch
import os

# Create visualizer
visualizer = build_visualizer(
    vis_dir="./visualizations/",
    img_dir="./data/FSC147_384_V2/images_384_VarV2/",
    activation="sigmoid",      # Apply sigmoid to outputs (None for raw)
    normalization=True,        # Normalize heatmap to [0, 1]
    with_image=True,           # Overlay on original image
)

# Visualize batch outputs
outputs = {
    "filename": ["image_001.jpg", "image_002.jpg"],
    "height": [768, 1024],
    "width": [1024, 768],
    "density": torch.randn(2, 1, 512, 512),       # Ground truth
    "density_pred": torch.randn(2, 1, 512, 512),  # Predictions
}
visualizer.vis_batch(outputs)
# Saves: ./visualizations/image_001_gt{X}_pred{Y}.png
#        ./visualizations/image_002_gt{X}_pred{Y}.png
```

--------------------------------

### Train, Evaluate, and Test Models via Shell Scripts

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Scripts for training, evaluating, and testing models using distributed data parallel support. Ensure you are in the correct experiment directory before execution.

```bash
# Navigate to experiment directory
cd experiments/FSC147/

# Train on 4 GPUs (GPU IDs: 0,1,2,3)
sh train_torch.sh 4 0,1,2,3

# Evaluate on validation set
sh eval_torch.sh 4 0,1,2,3

# Test on test set
sh test_torch.sh 4 0,1,2,3

# Alternative: Direct Python command
export PYTHONPATH=../../:$PYTHONPATH
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch \
    --nproc_per_node=4 \
    ../../tools/train_val.py \
    -c ./config.yaml

# Evaluation mode
CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch \
    --nproc_per_node=1 \
    ../../tools/train_val.py \
    -c ./config.yaml \
    -e  # Evaluate flag
```

--------------------------------

### UCSD Dataset Directory Structure

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

Set up the UCSD dataset directory with 'ucsdpeds_vidf', annotation files, and a mask. Ensure images are merged correctly.

```bash
|-- data
    |-- UCSD
        |-- ucsdpeds_vidf
        |-- gen_gt_density.py
        |-- train.json
        |-- test.json
        |-- exemplar.json
        |-- mask.png
```

--------------------------------

### Build Distributed DataLoaders

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Creates distributed data loaders for training, validation, and testing from a YAML configuration. Requires PyTorch distributed initialization and an EasyDict configuration object.

```python
import torch.distributed as dist
from datasets.data_builder import build_dataloader
from easydict import EasyDict

# Initialize distributed training
dist.init_process_group(backend="nccl")

# Dataset configuration
config = EasyDict({
    "type": "custom",                                      # or "custom_exemplar"
    "shot": 3,
    "img_dir": "./data/FSC147_384_V2/images_384_VarV2/",
    "density_dir": "./data/FSC147_384_V2/gt_density_map_adaptive_384_VarV2/",
    "input_size": [512, 512],
    "pixel_mean": [0.485, 0.456, 0.406],
    "pixel_std": [0.229, 0.224, 0.225],
    "batch_size": 1,
    "workers": 4,
    "train": {
        "meta_file": "./data/FSC147_384_V2/train.json",
        "hflip": {"prob": 0.5},
        "colorjitter": {"brightness": 0.2, "contrast": 0.2, "saturation": 0.2, "hue": 0.2, "prob": 0.5},
    },
    "val": {"meta_file": "./data/FSC147_384_V2/val.json"},
    "test": {"meta_file": "./data/FSC147_384_V2/test.json"},
})

train_loader, val_loader, test_loader = build_dataloader(config, distributed=True)

# Training loop
for batch in train_loader:
    images = batch["image"].cuda()       # [B, C, H, W]
    density = batch["density"].cuda()    # [B, 1, H, W]
    boxes = batch["boxes"].cuda()        # [B, M, 4]
    # ... forward pass
```

--------------------------------

### PUCPR+ Dataset Directory Structure

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

Arrange the PUCPR+ dataset files according to this structure. Include 'Images' and the necessary JSON annotation files.

```bash
|-- data
    |-- PUCPR+_devkit
        |-- Images
        |-- gen_gt_density.py
        |-- train.json
        |-- test.json
        |-- exemplar.json
```

--------------------------------

### Mall Dataset Directory Structure

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

Organize the Mall dataset files, including 'frames', annotation files, and 'mask.png'.

```bash
|-- data
    |-- Mall
        |-- frames
        |-- gen_gt_density.py
        |-- train.json
        |-- test.json
        |-- exemplar.json
        |-- mask.png
```

--------------------------------

### CARPK Dataset Directory Structure

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

Organize the CARPK dataset files under the specified directory structure. Ensure 'Images' and JSON annotation files are present.

```bash
|-- data
    |-- CARPK_devkit
        |-- Images
        |-- gen_gt_density.py
        |-- train.json
        |-- test.json
        |-- exemplar.json
```

--------------------------------

### Generate CARPK Ground-Truth Density Map

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

Execute this script to generate the ground-truth density maps for the CARPK dataset. Maps are saved to a dedicated subdirectory.

```python
python gen_gt_density.py
```

--------------------------------

### Train, Eval, or Test FSC147 (Slurm)

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

Use these commands to train, evaluate, or test the model on the FSC147 dataset when using a Slurm workload manager. Ensure you have the correct number of GPUs and partition specified.

```bash
sh train.sh #NUM_GPUS #PARTITION
```

```bash
sh eval.sh #NUM_GPUS #PARTITION
```

```bash
sh test.sh #NUM_GPUS #PARTITION
```

--------------------------------

### ShanghaiTech Dataset Directory Structure

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

Structure the ShanghaiTech dataset, separating 'part_A' and 'part_B' with their respective training/testing data, images, and annotation files.

```bash
|-- data
    |-- ShanghaiTech
        |-- part_A
            |-- train_data
                |-- images
            |-- test_data
                |-- images
            |-- gen_gt_density.py
            |-- train.json
            |-- test.json
            |-- exemplar.json
        |-- part_B
            |-- train_data
                |-- images
            |-- test_data
                |-- images
            |-- gen_gt_density.py
            |-- train.json
            |-- test.json
            |-- exemplar.json
```

--------------------------------

### Compute MAE and RMSE using Eval Helper

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Use `dump` to save predictions during evaluation and `merge_together` with `performances` to compute MAE and RMSE after all batches are processed. Alternatively, perform manual computation using NumPy.

```python
import numpy as np
from utils.eval_helper import dump, merge_together, performances

# During evaluation, dump predictions
outputs = {
    "filename": ["image_001.jpg"],
    "density": torch.randn(1, 1, 128, 128),      # Ground truth density
    "density_pred": torch.randn(1, 1, 128, 128), # Predicted density
}
dump(save_dir="./eval_results/", outputs=outputs)

# After all batches, merge and compute metrics
gt_cnts, pred_cnts = merge_together("./eval_results/")
mae, rmse = performances(gt_cnts, pred_cnts)
print(f"MAE: {mae:.2f}, RMSE: {rmse:.2f}")

# Manual computation
gt_counts = np.array([10.0, 25.0, 42.0, 15.0])
pred_counts = np.array([11.2, 23.5, 40.8, 16.1])
mae = np.mean(np.abs(pred_counts - gt_counts))
rmse = np.sqrt(np.mean((pred_counts - gt_counts) ** 2))
print(f"MAE: {mae:.2f}, RMSE: {rmse:.2f}")
# Output: MAE: 1.35, RMSE: 1.52
```

--------------------------------

### Build ResNet Backbone for Feature Extraction

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Use `build_backbone` to construct a ResNet model for feature extraction. Specify the `type`, `out_stride`, and `out_layers` to customize the backbone. An empty `pretrained_model` path defaults to ImageNet weights.

```python
from models.utils import build_backbone
import torch

# Build ResNet18 backbone
backbone = build_backbone(
    type="resnet18",                    # Options: resnet18, resnet50
    out_stride=4,                       # Output feature stride
    out_layers=[1, 2, 3],               # Layers to extract (1-4)
    pretrained_model="",                # Path to custom weights (empty = ImageNet)
)

print(f"Output dimension: {backbone.out_dim}")  # 448 for layers [1,2,3]

# Extract features
image = torch.randn(1, 3, 512, 512)
features = backbone(image)
print(f"Feature shape: {features.shape}")  # [1, 448, 128, 128] at stride 4
```

--------------------------------

### Generate Ground Truth Density Maps

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Creates a normalized density map from point annotations using a Gaussian filter.

```python
# Example: data/CARPK_devkit/gen_gt_density.py
import numpy as np
import cv2
from scipy.ndimage import gaussian_filter

def generate_density_map(image_shape, points, sigma=15):
    """
    Generate density map from point annotations.

    Args:
        image_shape: (height, width) of the image
        points: List of (x, y) point annotations
        sigma: Gaussian kernel sigma

    Returns:
        density_map: Normalized density map summing to len(points)
    """
    h, w = image_shape
    density = np.zeros((h, w), dtype=np.float32)

    for x, y in points:
        if 0 <= x < w and 0 <= y < h:
            density[int(y), int(x)] = 1.0

    density = gaussian_filter(density, sigma=sigma)

    # Normalize so sum equals point count
    if density.sum() > 0:
        density = density / density.sum() * len(points)

    return density

# Usage
points = [(100, 150), (200, 250), (300, 180)]  # Object centers
density_map = generate_density_map((512, 512), points, sigma=15)
print(f"Total count: {density_map.sum():.1f}")  # Should be ~3.0

# Save as numpy file
np.save("gt_density.npy", density_map)
```

--------------------------------

### Train or Eval FSC147 (Cross Validation - torch.distributed.launch)

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

Commands for training or evaluating the model in a cross-validation setting on FSC147 using torch.distributed.launch. Specify the total number of GPUs and their IDs. Note: Set config.saver.load_path for evaluation.

```bash
sh train_torch.sh #NUM_GPUS #GPU_IDS
```

```bash
sh eval_torch.sh #NUM_GPUS #GPU_IDS
```

--------------------------------

### Define Dataset Metadata Structures

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Defines annotation formats for FSC147 and CARPK-style datasets, including saving to JSON Lines.

```python
import json

# FSC147-style annotation (boxes in query image)
annotation_fsc147 = {
    "filename": "image_001.jpg",
    "density": "image_001.npy",
    "boxes": [
        [50, 80, 100, 150],     # [y_tl, x_tl, y_br, x_br]
        [200, 220, 260, 300],
        [120, 350, 180, 420],
    ]
}

# Save to JSON Lines format
with open("train.json", "w") as f:
    f.write(json.dumps(annotation_fsc147) + "\n")

# CARPK-style exemplar annotation (separate exemplar images)
annotation_exemplar = {
    "filename": "exemplar_001.jpg",
    "box": [100, 150, 200, 280],  # Single box in exemplar
}

# Query image annotation (no boxes needed)
annotation_query = {
    "filename": "query_image.jpg",
    "density": "query_image.npy",
}
```

--------------------------------

### Build SAFECount Exemplar Model for Class-Specific Counting

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Constructs the SAFECount model for class-specific counting using pre-extracted exemplar images. This variant accepts exemplar images and their corresponding bounding boxes instead of boxes within the query image. Configuration options include the number of SAFECount blocks, backbone type, and embedding dimensions.

```python
import torch
from models.safecount_exemplar import build_network

# Build model for exemplar-based counting
config = {
    "block": 2,
    "backbone": {
        "type": "resnet18",
        "out_layers": [1, 2, 3],
        "out_stride": 4,
    },
    "pool": {
        "type": "max",
        "size": [1, 1],
    },
    "embed_dim": 256,
    "mid_dim": 1024,
    "head": 8,
    "dropout": 0.0,
    "activation": "leaky_relu",
    "exemplar_scales": [],
    "initializer": {"method": "normal", "std": 0.001},
}

model = build_network(**config)
model.cuda()

# Input with separate exemplar images
sample = {
    "image": torch.randn(1, 3, 512, 512).cuda(),            # Query image
    "exemplar_imgs": torch.randn(1, 5, 3, 512, 512).cuda(), # 5 exemplar images [B, N, C, H, W]
    "exemplar_boxes": torch.tensor([[[100, 150, 200, 250],  # Box in each exemplar [B, N, 4]
                                      [80, 90, 160, 180],
                                      [50, 60, 130, 150],
                                      [90, 100, 180, 200],
                                      [70, 80, 150, 170]]]).float().cuda(),
}

with torch.no_grad():
    output = model(sample)
    count = output["density_pred"].sum().item()
    print(f"Predicted count: {count:.1f}")
```

--------------------------------

### Perform Inference with SAFECount

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Execute a forward pass using a sample dictionary containing image tensors and bounding box coordinates. The resulting density map is summed to calculate the final object count.

```python
sample = {
    "image": torch.randn(1, 3, 512, 512).cuda(),
    "boxes": torch.tensor([[[50, 80, 100, 150]]]).float().cuda(),
}
with torch.no_grad():
    output = model(sample)
    count = output["density_pred"].sum().item()
    print(f"Predicted count: {count:.1f}")
```

--------------------------------

### Extract Exemplar Features with crop_roi_feat

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Extracts feature regions corresponding to bounding boxes from feature maps using a specified stride.

```python
from models.utils import crop_roi_feat
import torch

# Feature map and bounding boxes
feat = torch.randn(1, 256, 128, 128)  # [1, C, H, W]
boxes = torch.tensor([
    [20, 30, 60, 80],    # [y_tl, x_tl, y_br, x_br] in image coordinates
    [100, 120, 150, 180],
]).float()

# Extract cropped features
out_stride = 4  # Feature map stride
feat_boxes = crop_roi_feat(feat, boxes, out_stride)

# Returns list of feature crops (variable sizes)
for i, crop in enumerate(feat_boxes):
    print(f"Box {i} feature shape: {crop.shape")
# Output: Box 0 feature shape: torch.Size([1, 256, 11, 13])
#         Box 1 feature shape: torch.Size([1, 256, 13, 16])
```

--------------------------------

### Build SAFECount Model for FSC147-style Counting

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Constructs the SAFECount model for few-shot counting where exemplar boxes are annotated within the query image. Requires a configuration dictionary specifying backbone, pooling, embedding dimensions, and attention heads. The model is then moved to CUDA for processing.

```python
import torch
from models.safecount import build_network

# Build SAFECount model for FSC147-style counting (exemplars from query image)
config = {
    "block": 4,                          # Number of SAFECount blocks
    "backbone": {
        "type": "resnet18",              # Options: resnet18, resnet50
        "out_layers": [1, 2, 3],         # Feature layers to use
        "out_stride": 4,                 # Output stride for feature maps
    },
    "pool": {
        "type": "max",                   # Pooling type: max or avg
        "size": [3, 3],                  # Pool size (must be odd)
    },
    "embed_dim": 256,                    # Embedding dimension
    "mid_dim": 1024,                     # Middle layer dimension
    "head": 8,                           # Number of attention heads
    "dropout": 0.0,                      # Dropout rate
    "activation": "leaky_relu",          # Activation: relu or leaky_relu
    "exemplar_scales": [0.8, 1, 1.2],    # Multi-scale exemplar factors
    "initializer": {
        "method": "normal",
        "std": 0.001,
    },
}

model = build_network(**config)
model.cuda()

# Prepare input (batch size must be 1)
sample = {
    "image": torch.randn(1, 3, 512, 512).cuda(),    # Query image [B, C, H, W]
    "boxes": torch.tensor([[[50, 80, 100, 150],     # Exemplar boxes [B, M, 4]
                            [200, 220, 250, 290]]]).float().cuda(),
}

# Forward pass returns density prediction
with torch.no_grad():
    output = model(sample)
    density_pred = output["density_pred"]
    predicted_count = density_pred.sum().item()     # Sum of density = object count
    print(f"Predicted count: {predicted_count:.1f}")
```

--------------------------------

### Train or Eval FSC147 (Cross Validation - Slurm)

Source: https://github.com/zhiyuanyou/safecount/blob/main/README.md

Commands for training or evaluating the model in a cross-validation setting on FSC147 using Slurm. Specify the number of GPUs and the partition.

```bash
sh train.sh #NUM_GPUS #PARTITION
```

```bash
sh eval.sh #NUM_GPUS #PARTITION
```

--------------------------------

### Build Density Regressor for Count Prediction

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

The `build_regressor` function creates a density regressor model. It takes the input feature dimension and activation function as parameters. The regressor converts feature maps to density maps.

```python
from models.utils import build_regressor
import torch

# Build regressor
regressor = build_regressor(
    in_dim=256,                 # Input feature dimension
    activation="leaky_relu",    # Activation function
)

# Feature map to density map (4x upsampling total)
features = torch.randn(1, 256, 128, 128)  # [B, C, H/4, W/4]
density = regressor(features)              # [B, 1, H, W]
print(f"Density shape: {density.shape}")   # [1, 1, 512, 512]
```

--------------------------------

### Calculate Predicted Object Count

Source: https://context7.com/zhiyuanyou/safecount/llms.txt

Computes the total object count from a density map tensor.

```python
predicted_count = density.sum().item()
print(f"Predicted count: {predicted_count:.1f}")
```
Tokens:
Raw