train.py

"""
Train YOLOv8 on the surgical instrument dataset.

Usage:
    python train.py
    python train.py --model yolov8m.pt --epochs 100 --batch 8 --device cpu
"""

import argparse
from pathlib import Path
from ultralytics import YOLO


def parse_args():
    p = argparse.ArgumentParser(description="Train YOLOv8 surgical instrument detector")
    p.add_argument("--model", default="yolov8s.pt",
                   help="Pretrained weights (n/s/m/l/x). Larger = more accurate, slower.")
    p.add_argument("--data", default="dataset/data.yaml")
    p.add_argument("--epochs", type=int, default=50)
    p.add_argument("--imgsz", type=int, default=640)
    p.add_argument("--batch", type=int, default=8)
    p.add_argument("--device", default="cpu",
                   help="GPU index (0, 1, ...) or 'cpu'")
    p.add_argument("--patience", type=int, default=15,
                   help="Early-stop patience (epochs without val improvement)")
    p.add_argument("--project", default="runs/train")
    p.add_argument("--name", default="surgical_s")
    return p.parse_args()


def _resolve_data_yaml(requested: str) -> Path:
    p = Path(requested)
    if p.exists():
        return p
    # Roboflow downloads put the yaml inside the dataset folder
    fallback = Path("dataset") / "data.yaml"
    if fallback.exists():
        print(f"data.yaml not found at '{requested}', using {fallback}")
        return fallback
    raise SystemExit(
        f"Dataset config not found: {requested}\n"
        "Run:  python dataset_setup.py --api-key YOUR_KEY"
    )


def main():
    args = parse_args()
    data_yaml = _resolve_data_yaml(args.data)

    model = YOLO(args.model)

    results = model.train(
        data=str(data_yaml),
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        device=args.device,
        patience=args.patience,
        project=args.project,
        name=args.name,
        exist_ok=True,
    )

    best_weights = Path(results.trainer.best)
    print(f"\nTraining complete.")
    print(f"Best weights : {best_weights}")
    print(f"Results dir  : {results.trainer.save_dir}")
    return results


if __name__ == "__main__":
    main()