Spaces:
Sleeping
Sleeping
File size: 3,057 Bytes
f120eab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """
VioLens Unified Model Training Script
Fine-tunes YOLOv8s on the Indian traffic dataset (19 classes)
with heavy data augmentation to compensate for small dataset size (284 images).
"""
import os
from ultralytics import YOLO
# Paths
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(SCRIPT_DIR)
DATA_YAML = os.path.join(PROJECT_ROOT, "dataset", "data.yaml")
MODELS_DIR = os.path.join(SCRIPT_DIR, "models")
OUTPUT_DIR = os.path.join(SCRIPT_DIR, "runs")
def train():
print("=" * 60)
print("VioLens Unified Model Training")
print("=" * 60)
print(f"Dataset config: {DATA_YAML}")
print(f"Output dir: {OUTPUT_DIR}")
print()
# Load pretrained YOLOv8s (small — better accuracy than nano)
model = YOLO("yolov8s.pt")
# Train with aggressive augmentation for small dataset
results = model.train(
data=DATA_YAML,
epochs=100,
imgsz=640,
batch=16,
name="violens_unified",
project=OUTPUT_DIR,
exist_ok=True,
# --- Data augmentation (aggressive for 284 images) ---
augment=True,
mosaic=1.0, # mosaic augmentation (combine 4 images)
mixup=0.3, # mix two images together
hsv_h=0.015, # hue jitter (handles lighting variation)
hsv_s=0.7, # saturation jitter (rain/fog/dust)
hsv_v=0.4, # value/brightness jitter
degrees=10.0, # rotation (slight camera angle variation)
translate=0.1, # translation
scale=0.5, # scale jitter
fliplr=0.5, # horizontal flip
flipud=0.0, # no vertical flip (vehicles don't appear upside down)
erasing=0.2, # random erasing (occlusion simulation)
crop_fraction=1.0, # crop fraction for classification (not used in detection)
# --- Training config ---
optimizer="AdamW",
lr0=0.001,
lrf=0.01, # final learning rate = lr0 * lrf
warmup_epochs=5,
weight_decay=0.0005,
patience=20, # early stopping if no improvement for 20 epochs
# --- GPU/CPU ---
device='cpu', # PyTorch didn't detect CUDA, using CPU
workers=4,
# --- Logging ---
verbose=True,
plots=True,
)
# Copy best weights to models dir
best_pt = os.path.join(OUTPUT_DIR, "violens_unified", "weights", "best.pt")
dest_pt = os.path.join(MODELS_DIR, "violens_unified.pt")
if os.path.exists(best_pt):
import shutil
os.makedirs(MODELS_DIR, exist_ok=True)
shutil.copy2(best_pt, dest_pt)
print(f"\n✅ Best model copied to: {dest_pt}")
# Verify
test_model = YOLO(dest_pt)
print(f"Model classes ({len(test_model.names)}): {test_model.names}")
else:
print(f"\n❌ Training did not produce best.pt at {best_pt}")
print("\n" + "=" * 60)
print("Training complete!")
print("=" * 60)
if __name__ == "__main__":
train()
|