Spaces:
Sleeping
Sleeping
File size: 1,664 Bytes
08d66df |
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 |
# ============================================================
# SMARTVISION AI - YOLOv8 TRAINING SCRIPT
# - Fine-tunes yolov8s on 25-class SmartVision detection dataset
# ============================================================
import os
import torch
from ultralytics import YOLO
# ------------------------------------------------------------
# 1. PATHS & CONFIG
# ------------------------------------------------------------
BASE_DIR = "smartvision_dataset"
DET_DIR = os.path.join(BASE_DIR, "detection")
DATA_YAML = os.path.join(DET_DIR, "data.yaml")
# YOLO model size:
# - yolov8n.pt : nano
# - yolov8s.pt : small (good tradeoff) ✅
MODEL_WEIGHTS = "yolov8s.pt"
# Auto-select device
device = "0" if torch.cuda.is_available() else "cpu"
print("🚀 Using device:", device)
print("📂 DATA_YAML:", DATA_YAML)
# ------------------------------------------------------------
# 2. LOAD BASE MODEL
# ------------------------------------------------------------
print(f"📥 Loading YOLOv8 model from: {MODEL_WEIGHTS}")
model = YOLO(MODEL_WEIGHTS)
# ------------------------------------------------------------
# 3. TRAIN
# ------------------------------------------------------------
results = model.train(
data=DATA_YAML,
epochs=50,
imgsz=640,
batch=8, # smaller for CPU
lr0=0.01,
optimizer="SGD",
device=device,
project="yolo_runs",
name="smartvision_yolov8s",
pretrained=True,
plots=True,
verbose=True,
)
print("\n✅ YOLO training complete.")
print("📁 Run directory: yolo_runs/smartvision_yolov8s/")
print("📦 Best weights: yolo_runs/smartvision_yolov8s/weights/best.pt")
|