DirtyLeague_Bot/yolo_detector.py

"""
DirtyLeague YOLOv8 ONNX Detector.

High-performance, pure NumPy + ONNXRuntime implementation of YOLOv8 object detection.
Independent of PyTorch runtime - lightweight and fast (~15-25ms per frame).
"""

import os
from dataclasses import dataclass
from typing import List, Optional, Tuple, Dict
import numpy as np
import cv2
import onnxruntime as ort


CLASSES = [
    "btn_fight",         # 0
    "btn_exit",          # 1
    "btn_leave",         # 2
    "btn_ok",            # 3
    "btn_collect",       # 4
    "btn_turn_all",      # 5
    "btn_close",         # 6
    "btn_remove_chest",  # 7
    "btn_open_chest",    # 8
    "crown_chest",       # 9
    "card",              # 10
]


@dataclass
class Detection:
    class_id: int
    class_name: str
    confidence: float
    box: Tuple[int, int, int, int]  # x1, y1, x2, y2 in original image pixels
    center: Tuple[int, int]         # cx, cy in original image pixels

    @property
    def width(self) -> int:
        return self.box[2] - self.box[0]

    @property
    def height(self) -> int:
        return self.box[3] - self.box[1]


class YoloDetector:
    def __init__(
        self,
        model_path: str = "models/dirty_league_yolo.onnx",
        conf_thres: float = 0.40,
        iou_thres: float = 0.45,
    ):
        if not os.path.exists(model_path):
            # Try finding relative to this file
            alt_path = os.path.join(os.path.dirname(__file__), model_path)
            if os.path.exists(alt_path):
                model_path = alt_path
            else:
                raise FileNotFoundError(f"Model not found at: {model_path}")

        # Configure ONNX session options
        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        sess_options.intra_op_num_threads = 4

        # Prefer CPU Execution Provider (reliable on all machines)
        providers = ["CPUExecutionProvider"]
        self.session = ort.InferenceSession(model_path, sess_options, providers=providers)

        # Get input specs
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape  # [1, 3, 640, 640]
        self.img_size = (self.input_shape[2], self.input_shape[3])

        self.conf_thres = conf_thres
        self.iou_thres = iou_thres

    def _letterbox(
        self, img: np.ndarray, target_shape=(640, 640)
    ) -> Tuple[np.ndarray, float, Tuple[int, int]]:
        """Resize image with aspect ratio preservation and padding."""
        h, w = img.shape[:2]
        tw, th = target_shape

        scale = min(tw / w, th / h)
        nw, nh = int(round(w * scale)), int(round(h * scale))
        resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)

        pad_w = (tw - nw) / 2
        pad_h = (th - nh) / 2
        top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
        left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))

        padded = cv2.copyMakeBorder(
            resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)
        )
        return padded, scale, (left, top)

    def _nms(self, boxes: np.ndarray, scores: np.ndarray) -> List[int]:
        """NumPy Non-Maximum Suppression."""
        x1 = boxes[:, 0]
        y1 = boxes[:, 1]
        x2 = boxes[:, 2]
        y2 = boxes[:, 3]

        areas = (x2 - x1) * (y2 - y1)
        order = scores.argsort()[::-1]

        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)
            if order.size == 1:
                break

            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            w = np.maximum(0.0, xx2 - xx1)
            h = np.maximum(0.0, yy2 - yy1)
            inter = w * h

            iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-7)
            inds = np.where(iou <= self.iou_thres)[0]
            order = order[inds + 1]

        return keep

    def detect(self, img_bgr: np.ndarray) -> List[Detection]:
        """
        Run detection on a BGR image.
        Returns list of Detection objects with coordinates in the original image space.
        """
        orig_h, orig_w = img_bgr.shape[:2]

        # 1. Letterbox resize
        padded, scale, (pad_x, pad_y) = self._letterbox(img_bgr, self.img_size)

        # 2. Preprocess: BGR -> RGB, HWC -> CHW, normalize 0..1
        rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB)
        blob = rgb.astype(np.float32) / 255.0
        blob = np.transpose(blob, (2, 0, 1))
        blob = np.expand_dims(blob, axis=0)

        # 3. ONNX inference
        outputs = self.session.run([self.output_name], {self.input_name: blob})[0]
        # Shape: [1, 15, 8400]
        predictions = outputs[0].T  # Transpose to [8400, 15]

        # Columns: cx, cy, w, h, class_0_score ... class_10_score
        boxes_xywh = predictions[:, :4]
        class_scores = predictions[:, 4:]

        # Find best class per anchor
        best_class_ids = np.argmax(class_scores, axis=1)
        confidences = np.max(class_scores, axis=1)

        # Filter by confidence threshold
        mask = confidences >= self.conf_thres
        boxes_xywh = boxes_xywh[mask]
        confidences = confidences[mask]
        best_class_ids = best_class_ids[mask]

        if len(boxes_xywh) == 0:
            return []

        # Convert cx, cy, w, h to x1, y1, x2, y2 in 640x640 space
        x1 = boxes_xywh[:, 0] - boxes_xywh[:, 2] / 2
        y1 = boxes_xywh[:, 1] - boxes_xywh[:, 3] / 2
        x2 = boxes_xywh[:, 0] + boxes_xywh[:, 2] / 2
        y2 = boxes_xywh[:, 1] + boxes_xywh[:, 3] / 2

        # Transform back to original image space (remove padding and scale)
        x1 = (x1 - pad_x) / scale
        y1 = (y1 - pad_y) / scale
        x2 = (x2 - pad_x) / scale
        y2 = (y2 - pad_y) / scale

        # Clip to original image boundaries
        x1 = np.clip(x1, 0, orig_w)
        y1 = np.clip(y1, 0, orig_h)
        x2 = np.clip(x2, 0, orig_w)
        y2 = np.clip(y2, 0, orig_h)

        boxes = np.stack([x1, y1, x2, y2], axis=1)

        # Per-class NMS
        detections: List[Detection] = []
        unique_classes = np.unique(best_class_ids)

        for c in unique_classes:
            cls_mask = best_class_ids == c
            cls_boxes = boxes[cls_mask]
            cls_confs = confidences[cls_mask]

            keep_idx = self._nms(cls_boxes, cls_confs)
            for idx in keep_idx:
                bx = cls_boxes[idx]
                x1_int, y1_int, x2_int, y2_int = int(bx[0]), int(bx[1]), int(bx[2]), int(bx[3])
                cx = (x1_int + x2_int) // 2
                cy = (y1_int + y2_int) // 2
                detections.append(
                    Detection(
                        class_id=int(c),
                        class_name=CLASSES[c] if c < len(CLASSES) else f"class_{c}",
                        confidence=float(cls_confs[idx]),
                        box=(x1_int, y1_int, x2_int, y2_int),
                        center=(cx, cy),
                    )
                )

        return detections

    def find_one(self, class_name: str, min_conf: Optional[float] = None) -> Optional[Detection]:
        """Find the single detection of class_name with highest confidence."""
        # Note: detect() needs to be called with an image; this is a helper on results
        pass

    def detect_dict(self, img_bgr: np.ndarray) -> Dict[str, List[Detection]]:
        """Convenience method returning detections grouped by class name."""
        results = self.detect(img_bgr)
        grouped: Dict[str, List[Detection]] = {}
        for det in results:
            grouped.setdefault(det.class_name, []).append(det)
        return grouped