""" 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