from ultralytics import YOLO
import os

# --- Configurations ---
DATA_YAML_PATH = 'data.yaml'
# Ye path confirm kar lein ke aapke folder structure ke mutabiq sahi ho.
# Maslan agar Helmet - No helmet detection.v1i.yolov8 folder C drive mein hai, 
# toh uska pura path de sakte hain: 'C:/Users/YourName/.../data.yaml'

EPOCHS = 50 # Behtar accuracy ke liye 50 epochs

# 1. Load a pretrained YOLOv8n model
model = YOLO('yolov8n.pt') 

# 2. Train the model
print(f"Starting training with data from: {DATA_YAML_PATH}")
if not os.path.exists(DATA_YAML_PATH):
    print(f"ERROR: data.yaml file not found at {DATA_YAML_PATH}")
    print("Please make sure the path is correct and the file exists.")
else:
    results = model.train(
        data=DATA_YAML_PATH,  
        epochs=EPOCHS,        
        imgsz=640,            
        plots=True            
    )
    print("Training completed! Check 'runs/detect/train/weights/best.pt' for your trained model.")