train.py

from __future__ import annotations

import argparse
import sys


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Train a YOLO11 sheep detector with recall-oriented defaults."
    )
    parser.add_argument("--data", default="dataset/sheep.yaml", help="Path to YOLO dataset yaml.")
    parser.add_argument("--model", default="yolo11n.pt", help="Base YOLO11 model or checkpoint path.")
    parser.add_argument("--epochs", type=int, default=100, help="Training epochs.")
    parser.add_argument("--imgsz", type=int, default=1280, help="Training image size.")
    parser.add_argument("--batch", default="16", help="Batch size, or auto values like -1.")
    parser.add_argument("--device", default=None, help="Training device, such as cpu, mps, or cuda:0.")
    parser.add_argument("--project", default="runs/sheep", help="Training output root directory.")
    parser.add_argument("--name", default="train", help="Experiment name under the project directory.")
    parser.add_argument("--workers", type=int, default=8, help="Number of dataloader workers.")
    parser.add_argument("--patience", type=int, default=30, help="Early stop patience.")
    parser.add_argument(
        "--amp",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Enable or disable Automatic Mixed Precision.",
    )
    parser.add_argument("--cache", action="store_true", help="Cache images for faster repeated training.")
    parser.add_argument(
        "--single-cls",
        action="store_true",
        help="Treat all labels as one class. Useful when the dataset contains only sheep.",
    )
    parser.add_argument("--close-mosaic", type=int, default=10, help="Disable mosaic in final epochs.")
    parser.add_argument("--degrees", type=float, default=0.0, help="Rotation augmentation.")
    parser.add_argument("--shear", type=float, default=0.0, help="Shear augmentation.")
    parser.add_argument("--perspective", type=float, default=0.0, help="Perspective augmentation.")
    parser.add_argument("--fliplr", type=float, default=0.5, help="Horizontal flip probability.")
    parser.add_argument("--mosaic", type=float, default=1.0, help="Mosaic augmentation strength.")
    parser.add_argument("--mixup", type=float, default=0.0, help="MixUp augmentation strength.")
    parser.add_argument("--copy-paste", type=float, default=0.0, help="Copy-paste augmentation strength.")
    return parser


def _parse_batch(raw_batch: str) -> int | float:
    lowered = raw_batch.strip().lower()
    if lowered == "-1":
        return -1
    try:
        return int(lowered)
    except ValueError:
        return float(lowered)


def run(args: argparse.Namespace) -> int:
    try:
        from ultralytics import YOLO
        from src.config import TrainingConfig
    except ImportError as exc:
        print(
            "Failed to import training dependencies. Install them with `pip install -r requirements.txt`.\n"
            f"Details: {exc}",
            file=sys.stderr,
        )
        return 1

    config = TrainingConfig(
        model_path=args.model,
        data=args.data,
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=_parse_batch(str(args.batch)),
        device=args.device,
        project=args.project,
        name=args.name,
        workers=args.workers,
        patience=args.patience,
        amp=args.amp,
        cache=args.cache,
        close_mosaic=args.close_mosaic,
        degrees=args.degrees,
        shear=args.shear,
        perspective=args.perspective,
        fliplr=args.fliplr,
        mosaic=args.mosaic,
        mixup=args.mixup,
        copy_paste=args.copy_paste,
        single_cls=args.single_cls,
    )

    try:
        model = YOLO(config.model_path)
        results = model.train(**config.train_kwargs())
    except Exception as exc:  # pragma: no cover - integration path
        print(f"Training failed: {exc}", file=sys.stderr)
        return 1

    save_dir = getattr(results, "save_dir", None)
    if save_dir:
        print(f"Training complete. Outputs saved to: {save_dir}")
        print(f"Best weights: {save_dir / 'weights' / 'best.pt'}")
    else:
        print("Training complete.")
    return 0


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    return run(args)


if __name__ == "__main__":
    raise SystemExit(main())