inference.py

"""Standalone YOLO image inference for Hugging Face model repositories.

The module works with uploaded local weights such as models/yolo11n.pt and
falls back to Ultralytics model names such as yolo11n.pt when no local model is
present.
"""

from __future__ import annotations

import argparse
import json
import os
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterable

from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO


DEFAULT_MODEL_CANDIDATES = (
    "models/best.pt",
    "models/yolo11n.pt",
    "models/yolov8n.pt",
    "yolo11n.pt",
)


@dataclass(slots=True)
class Detection:
    """Serializable object detection result."""

    class_id: int
    class_name: str
    confidence: float
    bbox: list[float]


def resolve_model_path(model_path: str | os.PathLike[str] | None = None) -> str:
    """Return the model path/name to load.

    Priority:
    1. Explicit function/CLI argument.
    2. MODEL_PATH environment variable.
    3. Existing files under models/.
    4. Ultralytics default model name, which downloads on first use.
    """

    explicit = str(model_path or "").strip()
    if explicit:
        return explicit

    env_model = os.getenv("MODEL_PATH", "").strip()
    if env_model:
        return env_model

    for candidate in DEFAULT_MODEL_CANDIDATES:
        if Path(candidate).exists():
            return candidate

    return DEFAULT_MODEL_CANDIDATES[-1]


def load_model(model_path: str | os.PathLike[str] | None = None) -> YOLO:
    """Load a YOLO model from a local file or Ultralytics model name."""

    return YOLO(resolve_model_path(model_path))


def _resolve_classes(model: YOLO, classes: str | Iterable[str | int] | None) -> list[int] | None:
    if classes is None:
        return None

    if isinstance(classes, str):
        items = [item.strip() for item in classes.split(",") if item.strip()]
    else:
        items = [str(item).strip() for item in classes if str(item).strip()]

    if not items:
        return None

    names = getattr(model, "names", {}) or {}
    name_to_id = {str(name).lower(): int(idx) for idx, name in names.items()}
    resolved: list[int] = []
    for item in items:
        if item.isdigit():
            resolved.append(int(item))
        else:
            class_id = name_to_id.get(item.lower())
            if class_id is not None:
                resolved.append(class_id)

    return resolved or None


def predict(
    image: str | os.PathLike[str] | Image.Image,
    model: YOLO | None = None,
    model_path: str | os.PathLike[str] | None = None,
    conf: float = 0.35,
    iou: float = 0.5,
    imgsz: int = 1280,
    classes: str | Iterable[str | int] | None = "person",
    device: str | None = None,
) -> tuple[list[Detection], Image.Image]:
    """Run YOLO detection on one image and return detections plus annotation."""

    loaded_model = model or load_model(model_path)
    pil_image = image if isinstance(image, Image.Image) else Image.open(image)
    pil_image = pil_image.convert("RGB")
    class_ids = _resolve_classes(loaded_model, classes)

    results = loaded_model.predict(
        source=pil_image,
        conf=conf,
        iou=iou,
        imgsz=imgsz,
        classes=class_ids,
        device=device,
        verbose=False,
    )

    detections: list[Detection] = []
    names = getattr(loaded_model, "names", {}) or {}

    for result in results:
        if result.boxes is None:
            continue

        boxes = result.boxes.xyxy.cpu().numpy()
        confidences = result.boxes.conf.cpu().numpy()
        class_indexes = result.boxes.cls.cpu().numpy().astype(int)
        for bbox, score, class_id in zip(boxes, confidences, class_indexes, strict=False):
            detections.append(
                Detection(
                    class_id=int(class_id),
                    class_name=str(names.get(int(class_id), class_id)),
                    confidence=round(float(score), 4),
                    bbox=[round(float(value), 2) for value in bbox],
                )
            )

    annotated = draw_detections(pil_image, detections)
    return detections, annotated


def draw_detections(image: Image.Image, detections: list[Detection]) -> Image.Image:
    """Draw boxes and confidence labels on an RGB image."""

    annotated = image.copy()
    draw = ImageDraw.Draw(annotated)
    font = ImageFont.load_default()

    for detection in detections:
        x1, y1, x2, y2 = detection.bbox
        color = _color_for_class(detection.class_id)
        label = f"{detection.class_name} {detection.confidence:.2f}"

        draw.rectangle((x1, y1, x2, y2), outline=color, width=3)
        text_bbox = draw.textbbox((x1, y1), label, font=font)
        text_width = text_bbox[2] - text_bbox[0]
        text_height = text_bbox[3] - text_bbox[1]
        label_y = max(0, y1 - text_height - 8)
        draw.rectangle((x1, label_y, x1 + text_width + 8, label_y + text_height + 6), fill=color)
        draw.text((x1 + 4, label_y + 3), label, fill=(255, 255, 255), font=font)

    return annotated


def _color_for_class(class_id: int) -> tuple[int, int, int]:
    palette = (
        (35, 100, 170),
        (61, 163, 93),
        (222, 122, 40),
        (153, 80, 160),
        (199, 62, 82),
    )
    return palette[class_id % len(palette)]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run YOLO detection on an image.")
    parser.add_argument("--image", required=True, type=Path, help="Input image path")
    parser.add_argument("--output", default=Path("examples/annotated_output.jpg"), type=Path)
    parser.add_argument("--json", default=None, type=Path, help="Optional JSON detections path")
    parser.add_argument("--model", default=None, help="Path/name of model weights")
    parser.add_argument("--conf", default=0.35, type=float, help="Confidence threshold")
    parser.add_argument("--iou", default=0.5, type=float, help="NMS IoU threshold")
    parser.add_argument("--imgsz", default=1280, type=int, help="Inference image size")
    parser.add_argument(
        "--classes",
        default="person",
        help="Comma-separated class names or IDs. Use empty string for all classes.",
    )
    parser.add_argument("--device", default=None, help="Device, for example cpu, 0, or cuda:0")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    class_filter = args.classes if args.classes.strip() else None
    detections, annotated = predict(
        image=args.image,
        model_path=args.model,
        conf=args.conf,
        iou=args.iou,
        imgsz=args.imgsz,
        classes=class_filter,
        device=args.device,
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    annotated.save(args.output)

    payload = [asdict(detection) for detection in detections]
    if args.json:
        args.json.parent.mkdir(parents=True, exist_ok=True)
        args.json.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    else:
        print(json.dumps(payload, indent=2))

    print(f"Annotated image saved to {args.output}")


if __name__ == "__main__":
    main()