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: QUASI_TOTAL_IOA: float = 0.90 SMALL_CONTAINED_IOA: float = 0.85 SMALL_RATIO_MAX: float = 0.50 SINGLE_PLAYER_HUE_PIVOT: float = 90.0 CORNER_INDICES = {0, 5, 24, 29} 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 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]] = {} global_player_feats: List[np.ndarray] = [] global_player_refs: List[Tuple[int, int]] = [] if process_objects: bbox_model_results = self.bbox_model.predict(batch_images) else: bbox_model_results = None if bbox_model_results is not None: for frame_idx_in_batch, detection in enumerate(bbox_model_results): if not hasattr(detection, "boxes") or detection.boxes is None: continue frame_id = offset + frame_idx_in_batch img_bgr = batch_images[frame_idx_in_batch] boxes: List[BoundingBox] = [] for box in detection.boxes.data: x1, y1, x2, y2, conf, cls_id = box.tolist() boxes.append( BoundingBox( x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2), cls_id=int(cls_id), conf=float(conf), ) ) 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) boxes = self.suppress_quasi_total_containment(boxes) boxes = self.suppress_small_contained(boxes) for i, bb in enumerate(boxes): if int(bb.cls_id) == 2: hs = self._hs_feature_from_roi(img_bgr, bb) global_player_feats.append(hs) global_player_refs.append((frame_id, i)) bboxes[frame_id] = boxes global_cluster_centers: Optional[np.ndarray] = None if len(global_player_feats) >= 2: feats = np.vstack(global_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) global_cluster_centers = centers for (frame_id, box_idx), lbl in zip(global_player_refs, labels): bb = bboxes[frame_id][box_idx] bb.cls_id = 6 if int(lbl) == 0 else 7 elif len(global_player_feats) == 1: hue, _ = global_player_feats[0] frame_id, box_idx = global_player_refs[0] bb = bboxes[frame_id][box_idx] bb.cls_id = ( 6 if float(hue) < self.SINGLE_PLAYER_HUE_PIVOT else 7 ) if process_objects: for frame_idx_in_batch, img_bgr in enumerate(batch_images): frame_id = offset + frame_idx_in_batch frame_boxes = bboxes.get(frame_id) if not frame_boxes: continue self._reclass_extra_goalkeepers( img_bgr, frame_boxes, global_cluster_centers ) 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