Skip to content

Latest commit

 

History

7 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 

Repository files navigation

Longtail Object Detection Project

This is an object detection project targeting long-tail distribution problems in traffic monitoring scenarios, implemented using YOLOv8 with a two-stage training strategy.

Dataset on HuggingFace

Project Overview

This project addresses the following challenges:

  • Severe Class Imbalance: Cars account for 70%, while rare classes (person, motorcycle, HOV) are severely underrepresented
  • Training from Scratch: No pretrained weights used, complete training for the specific domain
  • Two-Stage Strategy: General training first, followed by fine-tuning for rare classes
  • Rare Class Boosting: Automatic oversampling mechanism to balance class distribution

Dataset

The dataset is too large to include in this repository. It has been uploaded to HuggingFace Hub.

Download the dataset: HuggingFace Dataset

Download Scripts

Bash

git lfs install
git clone https://huggingface.co/datasets/grandmaeatsadumpling/my-images-dataset

Python

from huggingface_hub import hf_hub_download

# Example: download a single file from the dataset
hf_hub_download(repo_id="grandmaeatsadumpling/my-images-dataset", filename="example.png")

Directory Structure

Longtail-Object-Detection/
├── README.md                   # This file
└── code_NN6144048/
    ├── config_better.yaml      # Main configuration file
    ├── train_advanced.py       # Two-stage training main script
    ├── kaggle.py               # Kaggle submission file generator
    ├── requirements.txt        # Dependencies
    ├── utils/                  # Utility library
    │   ├── data_prep.py        # Dataset preparation
    │   ├── rare_class_boost.py # Rare class boosting
    │   ├── ensemble.py         # Ensemble inference
    │   └── inference.py        # Inference utilities
    └── output_better2/         # Training output directory
        ├── stage1_best.pt      # Stage 1 best model
        ├── best_model_advanced.pt # Final model
        └── ...

Environment Setup

System Requirements

  • OS: Linux (Ubuntu 20.04+ recommended)
  • Python: 3.8 or higher
  • GPU: NVIDIA GPU with 11GB+ VRAM (RTX 2080 Ti or higher)
  • CUDA: 11.8 or higher
  • Disk Space: At least 20GB available

Installation Steps

1. Create Python Virtual Environment

# Using conda (recommended)
conda create -n longtail python=3.9
conda activate longtail

# Or using venv
python3 -m venv venv
source venv/bin/activate  # Linux/Mac
# venv\Scripts\activate   # Windows

2. Install PyTorch

Install PyTorch according to your CUDA version:

# CUDA 11.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# CUDA 12.1
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# CPU version (for testing only, not recommended for training)
pip install torch torchvision torchaudio

Verify installation:

python -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'CUDA Available: {torch.cuda.is_available()}')"

3. Install Dependencies

cd code_NN6144048
pip install -r requirements.txt

4. Verify Installation

# Test YOLOv8 installation
python -c "from ultralytics import YOLO; print('Ultralytics YOLOv8 installed successfully')"

# Check GPU
python -c "import torch; print(f'Available GPUs: {torch.cuda.device_count()}')"

Data Preparation

Dataset Structure

Organize your data into the following structure:

/path/to/data/
├── train/
│   ├── img0001.png
│   ├── img0001.txt
│   ├── img0002.png
│   ├── img0002.txt
│   └── ...
└── test/
    ├── img0001.png
    ├── img0002.png
    └── ...

Annotation Format

Annotation files should be in YOLO format (.txt), one object per line:

class_id center_x center_y width height

Where:

  • class_id: 0=car, 1=hov, 2=person, 3=motorcycle
  • All coordinates are normalized values (0-1 range)

Configure Data Path

Edit the paths in config_better.yaml:

data:
  root_dir: "/path/to/data"
  train_dir: "train"
  test_dir: "test"
  output_dir: "./output_better2"

Model Training

Full Two-Stage Training (Recommended)

Run the complete two-stage training pipeline:

cd code_NN6144048
python train_advanced.py --stage all --config config_better.yaml

This will execute sequentially:

  1. Stage 1: General training (500 epochs)
  2. Stage 2: Rare class fine-tuning (150 epochs)
  3. Validation: Final model evaluation

Staged Training

Stage 1 Only

python train_advanced.py --stage 1 --config config_better.yaml

Output: output_better2/stage1_best.pt

Stage 2 Only (Requires Stage 1 completion)

python train_advanced.py --stage 2 --config config_better.yaml

If Stage 1 model is in a non-default location:

python train_advanced.py --stage 2 --config config_better.yaml --stage1_model /path/to/stage1_model.pt

Output: output_better2/best_model_advanced.pt

Training Parameter Adjustment

Main parameters can be adjusted in config_better.yaml:

training_stage1:
  epochs: 500              # Number of training epochs
  batch_size: 64           # Batch size (adjust based on GPU memory)
  lr0: 0.01                # Initial learning rate
  device: "1,2,3"          # GPU device IDs (comma-separated for multi-GPU)

  class_weights:           # Class weights
    car: 0.8
    hov: 2.5
    person: 5.0            # Increased weight for rare classes
    motorcycle: 4.0

Monitor Training

Training automatically generates:

  • Training curves: output_better2/stage1_train/results.png
  • Validation results: output_better2/validation/
  • Checkpoints: output_better2/stage1_train/weights/

Use TensorBoard for monitoring (if installed):

tensorboard --logdir output_better2

Prediction and Submission

Generate Kaggle Submission File

Use the trained model to predict on the test set:

python kaggle.py \
    --model output_better2/best_model_advanced.pt \
    --test-dir /path/to/test \
    --output submission.csv \
    --conf 0.00001 \
    --iou 0.7 \
    --img-size 640

Parameter Description

  • --model: Path to trained model (.pt file)
  • --test-dir: Test image directory
  • --output: Output CSV file path
  • --conf: Confidence threshold (default: 0.00001, lower values produce more detections)
  • --iou: IoU threshold for NMS (default: 0.7, higher values reduce duplicate boxes)
  • --img-size: Inference image size (default: 640)

Output Format

Generated CSV file format:

Image_ID,PredictionString
1,0.95 100.5 200.3 50.2 80.1 0 0.87 150.0 220.0 60.0 90.0 2
2,0.92 80.0 100.0 70.0 100.0 0
...

PredictionString format:

conf xmin ymin width height class_id [conf xmin ymin width height class_id ...]

Adjusting Thresholds for Performance

Increase Precision (reduce false positives):

python kaggle.py --model ... --conf 0.001 --iou 0.7

Increase Recall (detect more objects):

python kaggle.py --model ... --conf 0.00001 --iou 0.5

Balanced Settings:

python kaggle.py --model ... --conf 0.0001 --iou 0.6

Batch Processing Multiple Models

If you have multiple trained models, you can batch generate submission files:

# Using Stage 1 model
python kaggle.py \
    --model output_better2/stage1_best.pt \
    --test-dir /path/to/test \
    --output submission_stage1.csv

# Using Stage 2 model
python kaggle.py \
    --model output_better2/best_model_advanced.pt \
    --test-dir /path/to/test \
    --output submission_stage2.csv

Advanced Features

1. Rare Class Boosting

The project automatically enables rare class oversampling (configured in config_better.yaml):

rare_class_boost:
  enabled: true
  person_oversample: 8      # 8x oversampling for person class
  motorcycle_oversample: 6  # 6x oversampling for motorcycle class

This automatically creates a boosted dataset at yolo_dataset/images/train_boosted/

2. Custom Augmentation Strategy

Adjust augmentation parameters in config_better.yaml:

training_stage1:
  augment:
    mosaic: 1.0         # Mosaic augmentation probability
    mixup: 0.2          # Mixup augmentation probability
    copy_paste: 0.5     # Copy-Paste probability
    hsv_h: 0.03         # Hue variation
    degrees: 20.0       # Rotation angle
    scale: 0.8          # Scale range

3. Multi-GPU Training

Specify multiple GPUs in config_better.yaml:

training_stage1:
  device: "0,1,2"  # Use GPU 0, 1, 2

Or specify on command line:

CUDA_VISIBLE_DEVICES=0,1,2 python train_advanced.py --stage all

4. Resume Training

Resume training from checkpoint:

# Stage 1
python train_advanced.py --stage 1 --resume output_better2/stage1_train/weights/last.pt

# Stage 2
python train_advanced.py --stage 2 --resume output_better2/stage2_finetune/weights/last.pt

Troubleshooting

Issue 1: CUDA Out of Memory

Error Message: RuntimeError: CUDA out of memory

Solutions:

  1. Reduce batch size:

    training_stage1:
      batch_size: 32  # Reduce from 64 to 32
  2. Use a smaller model:

    model:
      name: "yolov8n"  # Use nano version
  3. Reduce image size:

    dataset:
      image_size: 512  # Reduce from 640 to 512

Issue 2: Training Too Slow

Solutions:

  1. Increase worker count:

    training_stage1:
      workers: 8  # Adjust based on CPU cores
  2. Enable mixed precision training (enabled by default):

    amp: true
  3. Use faster data loader:

    pip install --upgrade pillow-simd

Issue 3: Poor Model Performance

Solutions:

  1. Increase training epochs:

    training_stage1:
      epochs: 800  # Increase from 500 to 800
  2. Adjust class weights (for underperforming classes):

    training_stage1:
      class_weights:
        person: 8.0  # Increase further
  3. Enable Test Time Augmentation (TTA):

    inference:
      use_tta: true

Issue 4: Submission File Format Error

Validate submission file:

python -c "
import pandas as pd
df = pd.read_csv('submission.csv')
print('Columns:', df.columns.tolist())
print('First 5 rows:')
print(df.head())
print(f'Total rows: {len(df)}')
"

Checkpoints:

  • CSV must have two columns: Image_ID and PredictionString
  • Image_ID must be integer
  • PredictionString format: conf xmin ymin width height class_id ...

Issue 5: Stage 1 Model Not Found

Error Message: Stage 1 model not found

Solutions:

  1. Confirm Stage 1 training is complete
  2. Check model path:
    ls -lh output_better2/stage1_best.pt
  3. Manually specify path:
    python train_advanced.py --stage 2 --stage1_model output_better2/stage1_train/weights/best.pt

Performance Benchmark

Hardware Configuration

  • GPU: 3× NVIDIA RTX 2080 Ti (11GB)
  • CPU: Intel Xeon or AMD Ryzen
  • RAM: 32GB+

Training Time

  • Stage 1 (500 epochs): ~12-18 hours
  • Stage 2 (150 epochs): ~4-6 hours
  • Total Training Time: ~16-24 hours

Inference Speed

  • Single Image: ~20-30ms
  • Batch Inference (batch=16): ~300-400ms
  • FPS: 30-50 FPS (640×640 images)

Expected Model Performance

  • mAP50: >0.75
  • mAP50-95: >0.45
  • Model Size: ~22.5 MB

Project Files


References

Frameworks

Related Papers

  • YOLO Series: Real-Time Object Detection
  • Long-Tail Learning: Addressing Class Imbalance
  • Data Augmentation: Copy-Paste, Mosaic, Mixup

License

This project is for academic research and competition use only.


Last Updated: October 31, 2025

About

Long-tail object detection in aerial images.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages