from pathlib import Path from typing import List, Tuple, Dict import sys import os from numpy import ndarray from pydantic import BaseModel sys.path.append(os.path.dirname(os.path.abspath(__file__))) from ultralytics import YOLO from team_cluster import TeamClassifier from utils import ( BoundingBox, Constants, ) import time import torch import gc import cv2 import numpy as np from pitch import process_batch_input, get_cls_net import yaml 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: SMALL_CONTAINED_IOA = Constants.SMALL_CONTAINED_IOA SMALL_RATIO_MAX = Constants.SMALL_RATIO_MAX SINGLE_PLAYER_HUE_PIVOT = Constants.SINGLE_PLAYER_HUE_PIVOT CORNER_INDICES = Constants.CORNER_INDICES KEYPOINTS_CONFIDENCE = Constants.KEYPOINTS_CONFIDENCE CORNER_CONFIDENCE = Constants.CORNER_CONFIDENCE GOALKEEPER_POSITION_MARGIN = Constants.GOALKEEPER_POSITION_MARGIN MIN_SAMPLES_FOR_FIT = 16 # Minimum player crops needed before fitting TeamClassifier MAX_SAMPLES_FOR_FIT = 400 # Reduced for faster fitting (was 600) # Multi-scale detection parameters MULTI_SCALE_ENABLED = True MULTI_SCALE_SCALES = [1.0, 1.2, 0.8] # Original, larger, smaller MULTI_SCALE_NMS_THRESHOLD = 0.45 # NMS threshold for multi-scale def __init__(self, path_hf_repo: Path) -> None: try: device = "cuda" if torch.cuda.is_available() else "cpu" model_path = path_hf_repo / "football_object_detection.onnx" self.bbox_model = YOLO(model_path) print("BBox Model Loaded") team_model_path = path_hf_repo / "osnet_model.pth.tar-100" self.team_classifier = TeamClassifier( device=device, batch_size=32, model_name=str(team_model_path) ) print("Team Classifier Loaded") # Team classification state self.team_classifier_fitted = False self.player_crops_for_fit = [] model_kp_path = path_hf_repo / 'keypoint' config_kp_path = path_hf_repo / 'hrnetv2_w48.yaml' cfg_kp = yaml.safe_load(open(config_kp_path, 'r')) loaded_state_kp = torch.load(model_kp_path, map_location=device) model = get_cls_net(cfg_kp) model.load_state_dict(loaded_state_kp) model.to(device) model.eval() self.keypoints_model = model self.kp_threshold = 0.1 self.pitch_batch_size = 4 self.health = "healthy" print("✅ Keypoints Model Loaded") except Exception as e: self.health = "❌ Miner initialization failed: " + str(e) print(self.health) def __repr__(self) -> str: if self.health == 'healthy': return ( f"health: {self.health}\n" f"BBox Model: {type(self.bbox_model).__name__}\n" f"Keypoints Model: {type(self.keypoints_model).__name__}" ) else: return self.health def _calculate_iou(self, box1: Tuple[float, float, float, float], box2: Tuple[float, float, float, float]) -> float: """ Calculate Intersection over Union (IoU) between two bounding boxes. Args: box1: (x1, y1, x2, y2) box2: (x1, y1, x2, y2) Returns: IoU score (0-1) """ x1_1, y1_1, x2_1, y2_1 = box1 x1_2, y1_2, x2_2, y2_2 = box2 # Calculate intersection area x_left = max(x1_1, x1_2) y_top = max(y1_1, y1_2) x_right = min(x2_1, x2_2) y_bottom = min(y2_1, y2_2) if x_right < x_left or y_bottom < y_top: return 0.0 intersection_area = (x_right - x_left) * (y_bottom - y_top) # Calculate union area box1_area = (x2_1 - x1_1) * (y2_1 - y1_1) box2_area = (x2_2 - x1_2) * (y2_2 - y1_2) union_area = box1_area + box2_area - intersection_area if union_area == 0: return 0.0 return intersection_area / union_area def _multi_scale_detection_single_image(self, img_bgr: ndarray) -> List[Tuple[float, float, float, float, float, int]]: """ Multi-Scale Object Detection for improved small object detection. Uses multiple image scales and combines results with intelligent NMS. """ H, W = img_bgr.shape[:2] scales = self.MULTI_SCALE_SCALES all_detections = [] for scale in scales: if scale != 1.0: new_h, new_w = int(H * scale), int(W * scale) # Ensure dimensions are reasonable if new_h > 2048 or new_w > 2048 or new_h < 320 or new_w < 320: continue scaled_img = cv2.resize(img_bgr, (new_w, new_h), interpolation=cv2.INTER_LINEAR) else: scaled_img = img_bgr new_h, new_w = H, W # Run detection on scaled image results = self.bbox_model([scaled_img], verbose=False, save=False) if results and len(results) > 0 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 = 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 or y2 <= y1: continue # Boost confidence for detections at optimal scales box_area = (x2 - x1) * (y2 - y1) if scale == 1.2 and box_area < 2500: # Small objects benefit from upscaling conf *= 1.10 elif scale == 0.8 and box_area > 8000: # Large objects benefit from downscaling conf *= 1.05 all_detections.append((x1, y1, x2, y2, conf, cls_id)) # Apply NMS to remove duplicates if len(all_detections) == 0: return [] # Convert to numpy for NMS boxes_np = np.array([[d[0], d[1], d[2], d[3], d[4]] for d in all_detections]) cls_ids = np.array([d[5] for d in all_detections]) # Apply NMS per class keep_indices = [] for cls_id in np.unique(cls_ids): cls_mask = cls_ids == cls_id cls_boxes = boxes_np[cls_mask] if len(cls_boxes) == 0: continue # Calculate IoU matrix x1 = cls_boxes[:, 0] y1 = cls_boxes[:, 1] x2 = cls_boxes[:, 2] y2 = cls_boxes[:, 3] scores = cls_boxes[:, 4] areas = (x2 - x1) * (y2 - y1) order = scores.argsort()[::-1] keep = [] while len(order) > 0: i = order[0] keep.append(i) if len(order) == 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, xx2 - xx1) h = np.maximum(0, yy2 - yy1) intersection = w * h union = areas[i] + areas[order[1:]] - intersection iou = intersection / (union + 1e-6) inds = np.where(iou <= self.MULTI_SCALE_NMS_THRESHOLD)[0] order = order[inds + 1] # Map back to original indices cls_indices = np.where(cls_mask)[0] keep_indices.extend(cls_indices[keep].tolist()) # Return filtered detections filtered_detections = [] for idx in keep_indices: x1, y1, x2, y2, conf, cls_id = all_detections[idx] filtered_detections.append((x1, y1, x2, y2, conf, cls_id)) return filtered_detections def _detect_objects_batch(self, decoded_images: List[ndarray]) -> List: """ Detect objects in batch with optional multi-scale detection. Returns list of detection results compatible with original format. """ if self.MULTI_SCALE_ENABLED: # Multi-scale detection per image detection_results = [] for img in decoded_images: detections = self._multi_scale_detection_single_image(img) # Create a mock result object compatible with original format class MockBoxes: def __init__(self, detections): self.data = np.array([[d[0], d[1], d[2], d[3], d[4], d[5]] for d in detections], dtype=np.float32) class MockResult: def __init__(self, boxes): self.boxes = boxes mock_boxes = MockBoxes(detections) mock_result = MockResult(mock_boxes) detection_results.append(mock_result) return detection_results else: # Original single-scale detection batch_size = 16 detection_results = [] n_frames = len(decoded_images) for frame_number in range(0, n_frames, batch_size): batch_images = decoded_images[frame_number: frame_number + batch_size] detections = self.bbox_model(batch_images, verbose=False, save=False) detection_results.extend(detections) return detection_results def _team_classify(self, detection_results, decoded_images, offset): self.team_classifier_fitted = False start = time.time() # Collect player crops from first batch for fitting fit_sample_size = 500 # Reduced from 600 for faster fitting player_crops_for_fit = [] for frame_id in range(len(detection_results)): detection_box = detection_results[frame_id].boxes.data if len(detection_box) < 4: continue # Collect player boxes for team classification fitting (first batch only) if len(player_crops_for_fit) < fit_sample_size: frame_image = decoded_images[frame_id] for box in detection_box: x1, y1, x2, y2, conf, cls_id = box.tolist() if conf < 0.5: continue mapped_cls_id = str(int(cls_id)) # Only collect player crops (cls_id = 2) if mapped_cls_id == '2': crop = frame_image[int(y1):int(y2), int(x1):int(x2)] if crop.size > 0: player_crops_for_fit.append(crop) # Fit team classifier after collecting samples if self.team_classifier and not self.team_classifier_fitted and len(player_crops_for_fit) >= fit_sample_size: print(f"Fitting TeamClassifier with {len(player_crops_for_fit)} player crops") self.team_classifier.fit(player_crops_for_fit) self.team_classifier_fitted = True break if not self.team_classifier_fitted and len(player_crops_for_fit) >= 16: print(f"Fallback: Fitting TeamClassifier with {len(player_crops_for_fit)} player crops") self.team_classifier.fit(player_crops_for_fit) self.team_classifier_fitted = True end = time.time() print(f"Fitting Kmeans time: {end - start}") # Second pass: predict teams with configurable frame skipping optimization start = time.time() # Get configuration for frame skipping prediction_interval = 1 # Predict every frame (no skipping for better accuracy) iou_threshold = 0.35 # Increased from 0.3 for better matching print(f"Team classification - prediction_interval: {prediction_interval}, iou_threshold: {iou_threshold}") # Storage for predicted frame results: {frame_id: {box_idx: (bbox, team_id)}} predicted_frame_data = {} # Step 1: Predict for frames at prediction_interval only frames_to_predict = [] for frame_id in range(len(detection_results)): if frame_id % prediction_interval == 0: frames_to_predict.append(frame_id) print(f"Predicting teams for {len(frames_to_predict)}/{len(detection_results)} frames " f"(saving {100 - (len(frames_to_predict) * 100 // len(detection_results))}% compute)") for frame_id in frames_to_predict: detection_box = detection_results[frame_id].boxes.data frame_image = decoded_images[frame_id] # Collect player crops for this frame frame_player_crops = [] frame_player_indices = [] frame_player_boxes = [] for idx, box in enumerate(detection_box): x1, y1, x2, y2, conf, cls_id = box.tolist() if cls_id == 2 and conf < 0.55: # Lowered from 0.6 to catch more players continue mapped_cls_id = str(int(cls_id)) # Collect player crops for prediction if self.team_classifier and self.team_classifier_fitted and mapped_cls_id == '2': crop = frame_image[int(y1):int(y2), int(x1):int(x2)] if crop.size > 0: frame_player_crops.append(crop) frame_player_indices.append(idx) frame_player_boxes.append((x1, y1, x2, y2)) # Predict teams for all players in this frame if len(frame_player_crops) > 0: team_ids = self.team_classifier.predict(frame_player_crops) predicted_frame_data[frame_id] = {} for idx, bbox, team_id in zip(frame_player_indices, frame_player_boxes, team_ids): # Map team_id (0,1) to cls_id (6,7) team_cls_id = str(6 + int(team_id)) predicted_frame_data[frame_id][idx] = (bbox, team_cls_id) # Step 2: Process all frames (interpolate skipped frames) fallback_count = 0 interpolated_count = 0 bboxes: dict[int, list[BoundingBox]] = {} for frame_id in range(len(detection_results)): detection_box = detection_results[frame_id].boxes.data frame_image = decoded_images[frame_id] boxes = [] team_predictions = {} if frame_id % prediction_interval == 0: # Predicted frame: use pre-computed predictions if frame_id in predicted_frame_data: for idx, (bbox, team_cls_id) in predicted_frame_data[frame_id].items(): team_predictions[idx] = team_cls_id else: # Skipped frame: interpolate from neighboring predicted frames # Find nearest predicted frames prev_predicted_frame = (frame_id // prediction_interval) * prediction_interval next_predicted_frame = prev_predicted_frame + prediction_interval # Collect current frame player boxes for idx, box in enumerate(detection_box): x1, y1, x2, y2, conf, cls_id = box.tolist() if cls_id == 2 and conf < 0.55: # Lowered from 0.6 to catch more players continue mapped_cls_id = str(int(cls_id)) if self.team_classifier and self.team_classifier_fitted and mapped_cls_id == '2': target_box = (x1, y1, x2, y2) # Try to match with previous predicted frame best_team_id = None best_iou = 0.0 if prev_predicted_frame in predicted_frame_data: team_id, iou = self._find_best_match( target_box, predicted_frame_data[prev_predicted_frame], iou_threshold ) if team_id is not None: best_team_id = team_id best_iou = iou # Try to match with next predicted frame if available and no good match yet if best_team_id is None and next_predicted_frame < len(detection_results): if next_predicted_frame in predicted_frame_data: team_id, iou = self._find_best_match( target_box, predicted_frame_data[next_predicted_frame], iou_threshold ) if team_id is not None and iou > best_iou: best_team_id = team_id best_iou = iou # Track interpolation success if best_team_id is not None: interpolated_count += 1 else: # Fallback: if no match found, predict individually crop = frame_image[int(y1):int(y2), int(x1):int(x2)] if crop.size > 0: team_id = self.team_classifier.predict([crop])[0] best_team_id = str(6 + int(team_id)) fallback_count += 1 if best_team_id is not None: team_predictions[idx] = best_team_id # Parse boxes with team classification for idx, box in enumerate(detection_box): x1, y1, x2, y2, conf, cls_id = box.tolist() if cls_id == 2 and conf < 0.55: # Lowered from 0.6 to catch more players continue # Check overlap with staff box overlap_staff = False for idy, boxy in enumerate(detection_box): s_x1, s_y1, s_x2, s_y2, s_conf, s_cls_id = boxy.tolist() if cls_id == 2 and s_cls_id == 4: staff_iou = self._calculate_iou(box[:4], boxy[:4]) if staff_iou >= 0.8: overlap_staff = True break if overlap_staff: continue mapped_cls_id = str(int(cls_id)) # Override cls_id for players with team prediction if idx in team_predictions: mapped_cls_id = team_predictions[idx] if mapped_cls_id != '4': if int(mapped_cls_id) == 3 and conf < 0.5: continue boxes.append( BoundingBox( x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2), cls_id=int(mapped_cls_id), conf=float(conf), ) ) # Handle footballs - keep only the 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) bboxes[offset + frame_id] = boxes return bboxes def predict_batch(self, batch_images: List[ndarray], offset: int, n_keypoints: int) -> List[TVFrameResult]: start = time.time() detection_results = self._detect_objects_batch(batch_images) end = time.time() print(f"Detection time: {end - start}") start = time.time() bboxes = self._team_classify(detection_results, batch_images, offset) end = time.time() print(f"Team classify time: {end - start}") pitch_batch_size = min(self.pitch_batch_size, len(batch_images)) keypoints: Dict[int, List[Tuple[int, int]]] = {} start = time.time() while True: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() device_str = "cuda" keypoints_result = process_batch_input( batch_images, self.keypoints_model, self.kp_threshold, device_str, batch_size=pitch_batch_size, ) if keypoints_result is not None and len(keypoints_result) > 0: for frame_number_in_batch, kp_dict in enumerate(keypoints_result): if frame_number_in_batch >= len(batch_images): break frame_keypoints: List[Tuple[int, int]] = [] try: height, width = batch_images[frame_number_in_batch].shape[:2] if kp_dict is not None and isinstance(kp_dict, dict): for idx in range(32): x, y = 0, 0 kp_idx = idx + 1 if kp_idx in kp_dict: try: kp_data = kp_dict[kp_idx] if isinstance(kp_data, dict) and "x" in kp_data and "y" in kp_data: x = int(kp_data["x"] * width) y = int(kp_data["y"] * height) except (KeyError, TypeError, ValueError): pass frame_keypoints.append((x, y)) except (IndexError, ValueError, AttributeError): frame_keypoints = [(0, 0)] * 32 if len(frame_keypoints) < n_keypoints: frame_keypoints.extend([(0, 0)] * (n_keypoints - len(frame_keypoints))) else: frame_keypoints = frame_keypoints[:n_keypoints] keypoints[offset + frame_number_in_batch] = frame_keypoints break end = time.time() print(f"Keypoint time: {end - start}") results: List[TVFrameResult] = [] for frame_number in range(offset, offset + len(batch_images)): frame_boxes = bboxes.get(frame_number, []) frame_keypoints = keypoints.get(frame_number, [(0, 0) for _ in range(n_keypoints)]) result = TVFrameResult( frame_id=frame_number, boxes=frame_boxes, keypoints=frame_keypoints, ) results.append(result) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() return results