arm-model / model /README.md
pragadeeshv23's picture
Upload folder using huggingface_hub
5b86813 verified
# 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