# Road Anomaly Detection - Model Training

Complete training pipeline for road anomaly detection using YOLOv8, optimized for **NVIDIA RTX 2050 (4GB VRAM)** and **Raspberry Pi 4** deployment.

## 🚀 Quick Start

```fish
# 1. Setup environment
chmod +x setup_training.fish
./setup_training.fish

# 2. Activate environment
source .venv/bin/activate.fish

# 3. Verify your dataset
python verify_dataset.py

# 4. Train model
python train_road_anomaly_model.py

# 5. Package for Raspberry Pi
python package_for_rpi.py
```

## 📁 Project Structure

```
model/
├── setup_training.fish       # Environment setup (run first)
├── train_road_anomaly_model.py  # Main training script
├── verify_dataset.py         # Dataset validation
├── inference.py              # Run inference on images/video
├── package_for_rpi.py        # Create RPi deployment package
├── plan.txt                  # Detailed training plan
└── README.md                 # This file

../dataset/dataset/           # Your dataset (YOLO format)
├── data.yaml
├── train/images/  train/labels/
├── valid/images/  valid/labels/
└── test/images/   test/labels/
```

## 🎯 RTX 2050 Optimization

The training is pre-configured for 4GB VRAM:

| Parameter | Value | Reason |
|-----------|-------|--------|
| Batch Size | 8 | 4GB VRAM safe |
| Image Size | 416 | Edge-optimized |
| AMP | Enabled | 30-50% faster |
| Workers | 4 | 8-core CPU / 2 |
| Cache | Disabled | Save RAM for GPU |

### If you get OOM (Out of Memory):

Edit `train_road_anomaly_model.py`:
```python
BATCH_SIZE = 4   # Reduce from 8
IMAGE_SIZE = 320 # Reduce from 416
```

## 📊 Dataset Format

YOLO format with normalized coordinates:
```
<class_id> <x_center> <y_center> <width> <height>
```

Example label file:
```
0 0.5 0.5 0.2 0.15
1 0.3 0.7 0.1 0.08
```

Classes (default):
- 0: pothole
- 1: crack
- 2: bump
- 3: obstacle
- 4: road_damage

## 🔧 Scripts

### 1. Setup Environment
```fish
./setup_training.fish
```
Installs PyTorch with CUDA, Ultralytics, and all dependencies.

### 2. Verify Dataset
```bash
python verify_dataset.py [path/to/dataset]
```
Checks dataset structure, validates annotations, shows class distribution.

### 3. Train Model
```bash
python train_road_anomaly_model.py
```
Trains YOLOv8n model with RTX 2050 optimizations, exports to TFLite.

### 4. Run Inference
```bash
# Image
python inference.py --model road-anomaly-detection/train_*/weights/best.pt --source image.jpg

# Video
python inference.py --model best.pt --source video.mp4

# Camera
python inference.py --model best.pt --source camera

# TFLite model
python inference.py --model best.tflite --source image.jpg
```

### 5. Package for Raspberry Pi
```bash
python package_for_rpi.py
```
Creates deployment package with TFLite model and detection script.

## 📦 Output

After training:
```
road-anomaly-detection/
└── train_YYYYMMDD_HHMMSS/
    ├── weights/
    │   ├── best.pt              # PyTorch model
    │   └── best.tflite          # TFLite for RPi ⭐
    ├── training_report.txt      # Performance metrics
    ├── results.png              # Training curves
    ├── confusion_matrix.png     # Per-class accuracy
    └── sample_predictions/      # Test predictions
```

## 🍓 Raspberry Pi Deployment

```bash
# Transfer to RPi
scp -r rpi_deployment_*/ pi@raspberrypi.local:/home/pi/

# On RPi
cd rpi_deployment_*/
pip install -r requirements.txt
python detect.py --source 0  # Webcam
```

Expected performance:
- RPi 4: 10-12 FPS at 416x416
- RPi 5: 15-20 FPS

## 📈 Expected Results

| Metric | Target |
|--------|--------|
| mAP@50 | ≥ 0.80 |
| Precision | ≥ 0.75 |
| Recall | ≥ 0.75 |
| Model Size | ≤ 10 MB |
| RPi FPS | 10-12 |

## 🛠 Troubleshooting

### CUDA Out of Memory
```python
# In train_road_anomaly_model.py
BATCH_SIZE = 4
IMAGE_SIZE = 320
```

### Training Slow
```fish
# Check GPU utilization
nvidia-smi

# Should see 85-95% GPU usage
# If low, increase batch size
```

### No GPU Detected
```fish
# Check NVIDIA driver
nvidia-smi

# Check PyTorch CUDA
python -c "import torch; print(torch.cuda.is_available())"

# Reinstall PyTorch
uv pip install torch --index-url https://download.pytorch.org/whl/cu121 --force-reinstall
```

## 📚 References

- [Ultralytics YOLOv8](https://docs.ultralytics.com/)
- [TensorFlow Lite](https://www.tensorflow.org/lite)
- [PyTorch](https://pytorch.org/)

## 📄 License

MIT License