DronePeopleCounting

A PyTorch Lightning-based deep learning framework for automated people counting from drone imagery using density estimation techniques.

🚁 Overview

This project implements a computer vision pipeline for counting people in aerial drone footage. It uses density map regression to estimate crowd sizes by predicting density maps where each person is represented as a Gaussian blob. The framework is built on PyTorch Lightning for scalable and reproducible training.

Key Features

• Density Estimation: Uses Gaussian density maps for accurate people counting
• Multiple Architectures: Supports UNet and DeepLabV3Plus models
• Custom Loss Functions: Specialized counting losses (MAE, Relative MAE)
• Data Augmentation: Comprehensive augmentation pipeline using Albumentations
• Multi-GPU Training: Built-in support for distributed training
• Experiment Tracking: Neptune.ai integration for experiment management
• ONNX Export: Model export for deployment

📊 Supported Datasets

• DroneCrowd: Aerial imagery dataset for crowd counting
• CountingSim: Synthetic drone simulation dataset

🛠️ Installation

Requirements

• Python 3.9+
• CUDA-compatible GPU (recommended)

Setup

1. Clone the repository:

git clone https://github.com/PUTvision/DronePeopleCounting.git
cd DronePeopleCounting

2. Install dependencies:

pip install -r requirements.txt

3. Download datasets (optional):

# Download CountingSim dataset
python scripts/download_countingsim_dataset.py

4. Convert datasets to Zarr format for faster loading (optional):

# Convert CountingSim to Zarr
python scripts/generate_countingsim_zarr.py

🚀 Quick Start

Training

Train a model using the default DroneCrowd configuration:

python run.py

Train with custom configuration:

python run.py --config-name=config_dronecrowd.yaml

Configuration

The project uses Hydra for configuration management. Key configuration files:

• configs/config_dronecrowd.yaml: Main configuration for DroneCrowd dataset
• Model, data, and training parameters are configurable via YAML files

Key Configuration Options

# Model configuration
model:
  model_name: DeepLabV3Plus  # or UNet
  encoder_name: tu-semnasnet_075
  loss_function: MSE  # or MAE, CountingMAE, RelativeCountingMAE
  lr: 3e-4

# Data configuration  
datamodule:
  batch_size: 16
  image_size: [960, 544]
  augment: True

📁 Project Structure

DronePeopleCounting/
├── configs/                    # Hydra configuration files
│   └── config_dronecrowd.yaml
├── data/                       # Dataset storage
├── notebooks/                  # Jupyter notebooks
├── outputs/                    # Training outputs and logs
├── scripts/                    # Utility scripts
│   ├── download_*.py          # Dataset download scripts
│   └── generate_*_zarr.py     # Data conversion scripts
├── src/                        # Source code
│   ├── datamodules/           # Data loading and preprocessing
│   ├── losses/                # Custom loss functions
│   ├── metrics/               # Evaluation metrics
│   ├── models/                # Model architectures
│   └── utils/                 # Utilities
├── tests/                      # Unit tests
├── requirements.txt            # Dependencies
└── run.py                     # Main training script

🧠 Model Architecture

The framework supports multiple architectures:

1. DeepLabV3Plus: Encoder-decoder architecture with atrous convolutions
2. UNet: Classic U-shaped architecture for dense prediction

Density Estimation Approach

• Each person is represented as a Gaussian blob with value 400
• The model predicts density maps where pixel intensities represent person density
• Final count is obtained by integrating the density map and dividing by 400

📈 Loss Functions

• MSE/MAE: Standard regression losses on density maps
• CountingMAE: Mean Absolute Error on final count
• RelativeCountingMAE: Relative counting error for handling scale variations

🔧 Advanced Usage

Custom Datasets

To use your own dataset, implement a custom dataset class following the pattern in src/datamodules/datasets/:

class CustomDataset(Dataset):
    def __init__(self, data_root, images_list, image_size, sigma, augmentations):
        # Initialize dataset
        pass
    
    def __getitem__(self, index):
        # Load image and annotations
        # Apply augmentations
        # Generate density map
        return image, density_map

Multi-GPU Training

# Train on 2 GPUs with DDP
python run.py trainer.gpus=2 trainer.accelerator=ddp

Experiment Tracking

Set up Neptune.ai credentials:

export NEPTUNE_API_TOKEN="your_token"
export NEPTUNE_PROJECT_NAME="your_project"

Model Export

Export trained model to ONNX:

python run.py export.export_to_onnx=True

🧪 Testing

Run unit tests:

pytest tests/

Run specific test modules:

pytest tests/unit/test_losses.py
pytest tests/unit/test_metrics.py

📊 Evaluation Metrics

• Mean Absolute Error (MAE): Average absolute difference in counts
• Mean Squared Error (MSE): Average squared difference in counts
• Counting MAE: Task-specific counting error metric

🤝 Contributing

1. Fork the repository
2. Create a feature branch (git checkout -b feature/amazing-feature)
3. Commit your changes (git commit -m 'Add amazing feature')
4. Push to the branch (git push origin feature/amazing-feature)
5. Open a Pull Request