training/train_surface_gate_detector.py

from __future__ import annotations

import argparse
import shutil
from pathlib import Path

from ultralytics import YOLO


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Train a lightweight steel-surface detector and install it into the project."
    )
    parser.add_argument(
        "--data",
        default="training/data/surface_gate_detector/yolo_dataset/dataset.yaml",
        help="Path to the detector dataset YAML file.",
    )
    parser.add_argument(
        "--model",
        default="yolo11n.pt",
        help="Pretrained Ultralytics detector checkpoint to finetune.",
    )
    parser.add_argument("--epochs", type=int, default=18)
    parser.add_argument("--imgsz", type=int, default=512)
    parser.add_argument("--batch", type=float, default=16)
    parser.add_argument("--device", default="")
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--project", default="training/runs")
    parser.add_argument("--name", default="surface_gate_detector")
    parser.add_argument(
        "--target",
        default="models/steel_surface_detector.pt",
        help="Where to copy the best trained checkpoint for app inference.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    data_path = Path(args.data).resolve()
    if not data_path.exists():
        raise FileNotFoundError(f"Dataset YAML not found: {data_path}")

    batch = int(args.batch) if float(args.batch).is_integer() else float(args.batch)

    model = YOLO(args.model)
    results = model.train(
        data=str(data_path),
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=batch,
        device=args.device or None,
        workers=args.workers,
        project=args.project,
        name=args.name,
        optimizer="AdamW",
        patience=5,
        cos_lr=True,
        close_mosaic=4,
        hsv_h=0.01,
        hsv_s=0.15,
        hsv_v=0.15,
        translate=0.05,
        scale=0.15,
        fliplr=0.5,
        mosaic=0.3,
        mixup=0.0,
        cache=True,
        pretrained=True,
        single_cls=True,
        save=True,
        verbose=True,
    )

    best_path = Path(results.save_dir) / "weights" / "best.pt"
    if not best_path.exists():
        raise FileNotFoundError(f"Best checkpoint was not produced at {best_path}")

    target_path = Path(args.target).resolve()
    target_path.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(best_path, target_path)

    print(f"Training run saved to: {results.save_dir}")
    print(f"Best detector copied to: {target_path}")
    print("Next step: restart the FastAPI app so the steel ROI detector is used in the inspection pipeline.")


if __name__ == "__main__":
    main()