train.py

import os
from pathlib import Path

from ultralytics import YOLO
import torch


os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def main():
    base_dir = Path(__file__).resolve().parent
    model_path = base_dir / "models" / "yolov8n.pt"
    data_path = base_dir / "dataset" / "fruit_veg_v1" / "data.yaml"

    # Load the base model (local file prevents re-download)
    model = YOLO(model_path)

    # Kick off training
    model.train(
        data=str(data_path),
        epochs=50,
        imgsz=640,
        batch=16,  # RTX 4060 can try 16 or 32
        device=0,  # Use first GPU
        workers=0  # Windows stability
    )


if __name__ == "__main__":
    main()