detection/detector.py

"""YOLOv8 wrapper for ball, stump, and pad detection."""

from ultralytics import YOLO
from box import Box
from utils.io_utils import logger
import yaml
from pathlib import Path

# -------------------------------------------------
# Robust path to lbw_drs_ai/config/config.yaml
# -------------------------------------------------
ROOT = Path(__file__).resolve().parents[1]          # …/lbw_drs_ai
CFG_PATH = ROOT / "config" / "config.yaml"

if not CFG_PATH.is_file():
    raise FileNotFoundError(f"Config file not found: {CFG_PATH}")

CONFIG = Box(yaml.safe_load(CFG_PATH.read_text()))
# -------------------------------------------------

CLASS_MAP = {
    0: "ball",
    1: "stump",
    2: "pad",
}


class Detector:
    def __init__(self):
        logger.info("Loading YOLOv8 model…")
        self.model = YOLO(CONFIG.model.yolo_weights)
        self.conf = CONFIG.model.conf_threshold
        self.iou = CONFIG.model.iou_threshold

    def infer(self, frame):
        """Return list of detections as dicts: {cls, conf, bbox}."""
        results = self.model.predict(frame, conf=self.conf, iou=self.iou, verbose=False)
        detections = []
        for r in results:
            for b in r.boxes:
                cls_id = int(b.cls)
                detections.append({
                    "class": CLASS_MAP.get(cls_id, str(cls_id)),
                    "conf": float(b.conf),
                    "bbox": b.xyxy[0].cpu().numpy(),  # [x1,y1,x2,y2]
                })
        return detections