This is an object detection project targeting long-tail distribution problems in traffic monitoring scenarios, implemented using YOLOv8 with a two-stage training strategy.
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
The dataset is too large to include in this repository. It has been uploaded to HuggingFace Hub.
Download the dataset: HuggingFace Dataset
git lfs install
git clone https://huggingface.co/datasets/grandmaeatsadumpling/my-images-datasetfrom 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")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
└── ...
- 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
# 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 # WindowsInstall 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 torchaudioVerify installation:
python -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'CUDA Available: {torch.cuda.is_available()}')"cd code_NN6144048
pip install -r requirements.txt# 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()}')"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 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)
Edit the paths in config_better.yaml:
data:
root_dir: "/path/to/data"
train_dir: "train"
test_dir: "test"
output_dir: "./output_better2"Run the complete two-stage training pipeline:
cd code_NN6144048
python train_advanced.py --stage all --config config_better.yamlThis will execute sequentially:
- Stage 1: General training (500 epochs)
- Stage 2: Rare class fine-tuning (150 epochs)
- Validation: Final model evaluation
python train_advanced.py --stage 1 --config config_better.yamlOutput: output_better2/stage1_best.pt
python train_advanced.py --stage 2 --config config_better.yamlIf 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.ptOutput: output_better2/best_model_advanced.pt
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.0Training 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_better2Use 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--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)
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 ...]
Increase Precision (reduce false positives):
python kaggle.py --model ... --conf 0.001 --iou 0.7Increase Recall (detect more objects):
python kaggle.py --model ... --conf 0.00001 --iou 0.5Balanced Settings:
python kaggle.py --model ... --conf 0.0001 --iou 0.6If 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.csvThe 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 classThis automatically creates a boosted dataset at yolo_dataset/images/train_boosted/
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 rangeSpecify multiple GPUs in config_better.yaml:
training_stage1:
device: "0,1,2" # Use GPU 0, 1, 2Or specify on command line:
CUDA_VISIBLE_DEVICES=0,1,2 python train_advanced.py --stage allResume 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.ptError Message: RuntimeError: CUDA out of memory
Solutions:
-
Reduce batch size:
training_stage1: batch_size: 32 # Reduce from 64 to 32
-
Use a smaller model:
model: name: "yolov8n" # Use nano version
-
Reduce image size:
dataset: image_size: 512 # Reduce from 640 to 512
Solutions:
-
Increase worker count:
training_stage1: workers: 8 # Adjust based on CPU cores
-
Enable mixed precision training (enabled by default):
amp: true
-
Use faster data loader:
pip install --upgrade pillow-simd
Solutions:
-
Increase training epochs:
training_stage1: epochs: 800 # Increase from 500 to 800
-
Adjust class weights (for underperforming classes):
training_stage1: class_weights: person: 8.0 # Increase further
-
Enable Test Time Augmentation (TTA):
inference: use_tta: true
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_IDandPredictionString - Image_ID must be integer
- PredictionString format:
conf xmin ymin width height class_id ...
Error Message: Stage 1 model not found
Solutions:
- Confirm Stage 1 training is complete
- Check model path:
ls -lh output_better2/stage1_best.pt
- Manually specify path:
python train_advanced.py --stage 2 --stage1_model output_better2/stage1_train/weights/best.pt
- GPU: 3× NVIDIA RTX 2080 Ti (11GB)
- CPU: Intel Xeon or AMD Ryzen
- RAM: 32GB+
- Stage 1 (500 epochs): ~12-18 hours
- Stage 2 (150 epochs): ~4-6 hours
- Total Training Time: ~16-24 hours
- Single Image: ~20-30ms
- Batch Inference (batch=16): ~300-400ms
- FPS: 30-50 FPS (640×640 images)
- mAP50: >0.75
- mAP50-95: >0.45
- Model Size: ~22.5 MB
- Configuration File: config_better.yaml - All hyperparameter settings
- Training Script: train_advanced.py - Two-stage training implementation
- Prediction Script: kaggle.py - Kaggle submission generator
- YOLOv8: Ultralytics YOLOv8
- PyTorch: PyTorch Official
- YOLO Series: Real-Time Object Detection
- Long-Tail Learning: Addressing Class Imbalance
- Data Augmentation: Copy-Paste, Mosaic, Mixup
This project is for academic research and competition use only.
Last Updated: October 31, 2025