from pathlib import Path from typing import List, Tuple, Dict, Optional from ultralytics import YOLO from numpy import ndarray from pydantic import BaseModel import numpy as np import cv2 class BoundingBox(BaseModel): x1: int y1: int x2: int y2: int cls_id: int conf: float class TVFrameResult(BaseModel): frame_id: int boxes: List[BoundingBox] keypoints: List[Tuple[int, int]] class Miner: # Optimized for enumeration and placement - more aggressive detection QUASI_TOTAL_IOA: float = 0.88 # Slightly lower to keep more detections SMALL_CONTAINED_IOA: float = 0.82 # More lenient for small objects SMALL_RATIO_MAX: float = 0.55 # Allow slightly larger size differences SINGLE_PLAYER_HUE_PIVOT: float = 90.0 CORNER_INDICES = {0, 5, 24, 29} # Enumeration-specific constants AGGRESSIVE_SCALES = [1.0, 1.3, 0.7, 1.1, 0.9] # More scales for better coverage ENUMERATION_NMS_THRESHOLD = 0.4 # Lower NMS for better enumeration SMALL_OBJECT_CONF_BOOST = 1.15 # Boost confidence for small objects def __init__(self, path_hf_repo: Path) -> None: self.bbox_model = YOLO(path_hf_repo / "objdetect.pt") print("BBox Model (objdetect.pt) Loaded") self.keypoints_model = YOLO(path_hf_repo / "keypointdetect.pt") print("Keypoints Model (keypointdetect.pt) Loaded") def __repr__(self) -> str: return ( f"BBox Model: {type(self.bbox_model).__name__}\n" f"Keypoints Model: {type(self.keypoints_model).__name__}" ) @staticmethod def _clip_box_to_image(x1: int, y1: int, x2: int, y2: int, w: int, h: int) -> Tuple[int, int, int, int]: x1 = max(0, min(int(x1), w - 1)) y1 = max(0, min(int(y1), h - 1)) x2 = max(0, min(int(x2), w - 1)) y2 = max(0, min(int(y2), h - 1)) if x2 <= x1: x2 = min(w - 1, x1 + 1) if y2 <= y1: y2 = min(h - 1, y1 + 1) return x1, y1, x2, y2 @staticmethod def _area(bb: BoundingBox) -> int: return max(0, bb.x2 - bb.x1) * max(0, bb.y2 - bb.y1) @staticmethod def _intersect_area(a: BoundingBox, b: BoundingBox) -> int: ix1 = max(a.x1, b.x1) iy1 = max(a.y1, b.y1) ix2 = min(a.x2, b.x2) iy2 = min(a.y2, b.y2) if ix2 <= ix1 or iy2 <= iy1: return 0 return (ix2 - ix1) * (iy2 - iy1) @staticmethod def _center(bb: BoundingBox) -> Tuple[float, float]: return (0.5 * (bb.x1 + bb.x2), 0.5 * (bb.y1 + bb.y2)) @staticmethod def _mean_hs(img_bgr: np.ndarray) -> Tuple[float, float]: hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV) return float(np.mean(hsv[:, :, 0])), float(np.mean(hsv[:, :, 1])) def _hs_feature_from_roi(self, img_bgr: np.ndarray, box: BoundingBox) -> np.ndarray: H, W = img_bgr.shape[:2] x1, y1, x2, y2 = self._clip_box_to_image(box.x1, box.y1, box.x2, box.y2, W, H) roi = img_bgr[y1:y2, x1:x2] if roi.size == 0: return np.array([0.0, 0.0], dtype=np.float32) hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV) lower_green = np.array([35, 60, 60], dtype=np.uint8) upper_green = np.array([85, 255, 255], dtype=np.uint8) green_mask = cv2.inRange(hsv, lower_green, upper_green) non_green_mask = cv2.bitwise_not(green_mask) num_non_green = int(np.count_nonzero(non_green_mask)) total = hsv.shape[0] * hsv.shape[1] if num_non_green > max(50, total // 20): h_vals = hsv[:, :, 0][non_green_mask > 0] s_vals = hsv[:, :, 1][non_green_mask > 0] h_mean = float(np.mean(h_vals)) if h_vals.size else 0.0 s_mean = float(np.mean(s_vals)) if s_vals.size else 0.0 else: h_mean, s_mean = self._mean_hs(roi) return np.array([h_mean, s_mean], dtype=np.float32) def _ioa(self, a: BoundingBox, b: BoundingBox) -> float: inter = self._intersect_area(a, b) aa = self._area(a) if aa <= 0: return 0.0 return inter / aa def suppress_quasi_total_containment(self, boxes: List[BoundingBox]) -> List[BoundingBox]: if len(boxes) <= 1: return boxes keep = [True] * len(boxes) for i in range(len(boxes)): if not keep[i]: continue for j in range(len(boxes)): if i == j or not keep[j]: continue ioa_i_in_j = self._ioa(boxes[i], boxes[j]) if ioa_i_in_j >= self.QUASI_TOTAL_IOA: keep[i] = False break return [bb for bb, k in zip(boxes, keep) if k] def suppress_small_contained(self, boxes: List[BoundingBox]) -> List[BoundingBox]: if len(boxes) <= 1: return boxes keep = [True] * len(boxes) areas = [self._area(bb) for bb in boxes] for i in range(len(boxes)): if not keep[i]: continue for j in range(len(boxes)): if i == j or not keep[j]: continue ai, aj = areas[i], areas[j] if ai == 0 or aj == 0: continue if ai <= aj: ratio = ai / aj if ratio <= self.SMALL_RATIO_MAX: ioa_i_in_j = self._ioa(boxes[i], boxes[j]) if ioa_i_in_j >= self.SMALL_CONTAINED_IOA: keep[i] = False break else: ratio = aj / ai if ratio <= self.SMALL_RATIO_MAX: ioa_j_in_i = self._ioa(boxes[j], boxes[i]) if ioa_j_in_i >= self.SMALL_CONTAINED_IOA: keep[j] = False return [bb for bb, k in zip(boxes, keep) if k] def _assign_players_two_clusters(self, features: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 20, 1.0) _, labels, centers = cv2.kmeans( np.float32(features), K=2, bestLabels=None, criteria=criteria, attempts=5, flags=cv2.KMEANS_PP_CENTERS, ) return labels.reshape(-1), centers def _reclass_extra_goalkeepers( self, img_bgr: np.ndarray, boxes: List[BoundingBox], cluster_centers: Optional[np.ndarray], ) -> None: gk_idxs = [i for i, bb in enumerate(boxes) if int(bb.cls_id) == 1] if len(gk_idxs) <= 1: return gk_idxs_sorted = sorted(gk_idxs, key=lambda i: boxes[i].conf, reverse=True) keep_gk_idx = gk_idxs_sorted[0] to_reclass = gk_idxs_sorted[1:] for gki in to_reclass: hs_gk = self._hs_feature_from_roi(img_bgr, boxes[gki]) if cluster_centers is not None: d0 = float(np.linalg.norm(hs_gk - cluster_centers[0])) d1 = float(np.linalg.norm(hs_gk - cluster_centers[1])) assign_cls = 6 if d0 <= d1 else 7 else: assign_cls = 6 if float(hs_gk[0]) < self.SINGLE_PLAYER_HUE_PIVOT else 7 boxes[gki].cls_id = int(assign_cls) def _aggressive_multi_scale_detection(self, img_bgr: np.ndarray) -> List[BoundingBox]: """ Aggressive Multi-Scale Object Detection optimized for enumeration and placement. Uses 5 scales with confidence boosting for small objects. """ H, W = img_bgr.shape[:2] all_detections = [] for scale in self.AGGRESSIVE_SCALES: if scale != 1.0: new_h, new_w = int(H * scale), int(W * scale) # More lenient dimension constraints for aggressive detection if new_h > 2560 or new_w > 2560 or new_h < 256 or new_w < 256: continue scaled_img = cv2.resize(img_bgr, (new_w, new_h)) else: scaled_img = img_bgr new_h, new_w = H, W # Run detection on scaled image results = self.bbox_model.predict([scaled_img], verbose=False) if results and hasattr(results[0], "boxes") and results[0].boxes is not None: for box in results[0].boxes.data: x1, y1, x2, y2, conf, cls_id = box.tolist() # Scale coordinates back to original image size if scale != 1.0: x1 = x1 / scale y1 = y1 / scale x2 = x2 / scale y2 = y2 / scale # Clip to original image bounds x1, y1, x2, y2 = self._clip_box_to_image(x1, y1, x2, y2, W, H) # Calculate box area for confidence boosting box_area = (x2 - x1) * (y2 - y1) # Aggressive confidence boosting based on scale and size if scale == 1.3 and box_area < 1500: # Very small objects at high scale conf *= self.SMALL_OBJECT_CONF_BOOST elif scale == 1.1 and box_area < 3000: # Small objects at medium scale conf *= 1.10 elif scale == 0.7 and box_area > 15000: # Large objects at small scale conf *= 1.08 elif scale == 0.9 and box_area > 8000: # Medium-large objects conf *= 1.05 # Extra boost for small objects regardless of scale if box_area < 1000: conf *= 1.12 all_detections.append(BoundingBox( x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2), cls_id=int(cls_id), conf=float(conf) )) # Apply enumeration-optimized NMS return self._enumeration_optimized_nms(all_detections) def _enumeration_optimized_nms(self, boxes: List[BoundingBox]) -> List[BoundingBox]: """ Enumeration-optimized NMS with lower threshold to preserve more detections. """ if not boxes: return [] # Group by class for class-specific NMS boxes_by_class = {} for box in boxes: if box.cls_id not in boxes_by_class: boxes_by_class[box.cls_id] = [] boxes_by_class[box.cls_id].append(box) final_boxes = [] for cls_id, class_boxes in boxes_by_class.items(): # Sort by confidence class_boxes_sorted = sorted(class_boxes, key=lambda x: x.conf, reverse=True) keep = [] while class_boxes_sorted: # Take the highest confidence box current = class_boxes_sorted.pop(0) keep.append(current) # Remove boxes with high IoU (lower threshold for enumeration) remaining = [] for box in class_boxes_sorted: iou = self._calculate_iou(current, box) if iou < self.ENUMERATION_NMS_THRESHOLD: remaining.append(box) elif box.conf > current.conf * 0.95: # Keep very close confidence boxes remaining.append(box) class_boxes_sorted = remaining final_boxes.extend(keep) return final_boxes def _calculate_iou(self, box1: BoundingBox, box2: BoundingBox) -> float: """Calculate Intersection over Union (IoU) between two bounding boxes.""" # Calculate intersection x1 = max(box1.x1, box2.x1) y1 = max(box1.y1, box2.y1) x2 = min(box1.x2, box2.x2) y2 = min(box1.y2, box2.y2) if x2 <= x1 or y2 <= y1: return 0.0 intersection = (x2 - x1) * (y2 - y1) # Calculate union area1 = (box1.x2 - box1.x1) * (box1.y2 - box1.y1) area2 = (box2.x2 - box2.x1) * (box2.y2 - box2.y1) union = area1 + area2 - intersection return intersection / union if union > 0 else 0.0 def predict_batch( self, batch_images: List[ndarray], offset: int, n_keypoints: int, task_type: Optional[str] = None, ) -> List[TVFrameResult]: process_objects = task_type is None or task_type == "object" process_keypoints = task_type is None or task_type == "keypoint" bboxes: Dict[int, List[BoundingBox]] = {} if process_objects: # Use aggressive multi-scale detection for optimal enumeration and placement for frame_idx_in_batch, img_bgr in enumerate(batch_images): boxes = self._aggressive_multi_scale_detection(img_bgr) # Handle multiple football detections (keep best one) footballs = [bb for bb in boxes if int(bb.cls_id) == 0] if len(footballs) > 1: best_ball = max(footballs, key=lambda b: b.conf) boxes = [bb for bb in boxes if int(bb.cls_id) != 0] boxes.append(best_ball) # Apply more lenient suppression for better enumeration boxes = self.suppress_quasi_total_containment(boxes) boxes = self.suppress_small_contained(boxes) # Team classification for players player_indices: List[int] = [] player_feats: List[np.ndarray] = [] for i, bb in enumerate(boxes): if int(bb.cls_id) == 2: hs = self._hs_feature_from_roi(img_bgr, bb) player_indices.append(i) player_feats.append(hs) cluster_centers: Optional[np.ndarray] = None n_players = len(player_feats) if n_players >= 2: feats = np.vstack(player_feats) labels, centers = self._assign_players_two_clusters(feats) order = np.argsort(centers[:, 0]) centers = centers[order] remap = {old_idx: new_idx for new_idx, old_idx in enumerate(order)} labels = np.vectorize(remap.get)(labels) cluster_centers = centers for idx_in_list, lbl in zip(player_indices, labels): boxes[idx_in_list].cls_id = 6 if int(lbl) == 0 else 7 elif n_players == 1: hue, _ = player_feats[0] boxes[player_indices[0]].cls_id = 6 if float(hue) < self.SINGLE_PLAYER_HUE_PIVOT else 7 self._reclass_extra_goalkeepers(img_bgr, boxes, cluster_centers) bboxes[offset + frame_idx_in_batch] = boxes keypoints: Dict[int, List[Tuple[int, int]]] = {} if process_keypoints: keypoints_model_results = self.keypoints_model.predict(batch_images) else: keypoints_model_results = None if keypoints_model_results is not None: for frame_idx_in_batch, detection in enumerate(keypoints_model_results): if not hasattr(detection, "keypoints") or detection.keypoints is None: continue frame_keypoints_with_conf: List[Tuple[int, int, float]] = [] for i, part_points in enumerate(detection.keypoints.data): for k_id, (x, y, _) in enumerate(part_points): confidence = float(detection.keypoints.conf[i][k_id]) frame_keypoints_with_conf.append((int(x), int(y), confidence)) if len(frame_keypoints_with_conf) < n_keypoints: frame_keypoints_with_conf.extend( [(0, 0, 0.0)] * (n_keypoints - len(frame_keypoints_with_conf)) ) else: frame_keypoints_with_conf = frame_keypoints_with_conf[:n_keypoints] filtered_keypoints: List[Tuple[int, int]] = [] for idx, (x, y, confidence) in enumerate(frame_keypoints_with_conf): if idx in self.CORNER_INDICES: if confidence < 0.3: filtered_keypoints.append((0, 0)) else: filtered_keypoints.append((int(x), int(y))) else: if confidence < 0.5: filtered_keypoints.append((0, 0)) else: filtered_keypoints.append((int(x), int(y))) keypoints[offset + frame_idx_in_batch] = filtered_keypoints results: List[TVFrameResult] = [] for frame_number in range(offset, offset + len(batch_images)): results.append( TVFrameResult( frame_id=frame_number, boxes=bboxes.get(frame_number, []), keypoints=keypoints.get( frame_number, [(0, 0) for _ in range(n_keypoints)], ), ) ) return results