from pathlib import Path from typing import List, Tuple, Dict, Optional 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, classify_teams_batch, ) import time import torch import gc import cv2 import numpy as np from collections import defaultdict 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 = 1000 # Maximum samples to avoid overfitting 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(f"BBox Model Loaded: class name {self.bbox_model.names}") 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 = [] # self.keypoints_model = YOLO(path_hf_repo / "keypoint.pt") 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 print("Keypoints Model (keypoint.pt) Loaded") 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 _extract_jersey_region(self, crop: ndarray) -> ndarray: """ Extract jersey region (upper body) from player crop. For close-ups, focuses on upper 60%, for distant shots uses full crop. """ if crop is None or crop.size == 0: return crop h, w = crop.shape[:2] if h < 10 or w < 10: return crop # For close-up shots, extract upper body (jersey region) is_closeup = h > 100 or (h * w) > 12000 if is_closeup: # Upper 60% of the crop (jersey area, avoiding shorts) jersey_top = 0 jersey_bottom = int(h * 0.60) jersey_left = max(0, int(w * 0.05)) jersey_right = min(w, int(w * 0.95)) return crop[jersey_top:jersey_bottom, jersey_left:jersey_right] return crop def _extract_color_signature(self, crop: ndarray) -> Optional[np.ndarray]: """ Extract color signature from jersey region using HSV and LAB color spaces. Returns a feature vector with dominant colors and color statistics. """ if crop is None or crop.size == 0: return None jersey_region = self._extract_jersey_region(crop) if jersey_region.size == 0: return None try: # Convert to HSV and LAB color spaces hsv = cv2.cvtColor(jersey_region, cv2.COLOR_BGR2HSV) lab = cv2.cvtColor(jersey_region, cv2.COLOR_BGR2LAB) # Reshape for processing hsv_flat = hsv.reshape(-1, 3).astype(np.float32) lab_flat = lab.reshape(-1, 3).astype(np.float32) # Compute statistics for HSV hsv_mean = np.mean(hsv_flat, axis=0) / 255.0 hsv_std = np.std(hsv_flat, axis=0) / 255.0 # Compute statistics for LAB lab_mean = np.mean(lab_flat, axis=0) / 255.0 lab_std = np.std(lab_flat, axis=0) / 255.0 # Dominant color (most frequent hue) hue_hist, _ = np.histogram(hsv_flat[:, 0], bins=36, range=(0, 180)) dominant_hue = np.argmax(hue_hist) * 5 # Convert to hue value # Combine features color_features = np.concatenate([ hsv_mean, hsv_std, lab_mean[:2], # L and A channels (B is less informative) lab_std[:2], [dominant_hue / 180.0] # Normalized dominant hue ]) return color_features except Exception as e: print(f"Error extracting color signature: {e}") return None def _get_spatial_position(self, bbox: Tuple[float, float, float, float], frame_width: int, frame_height: int) -> Tuple[float, float]: """ Get normalized spatial position of player on the pitch. Returns (x_normalized, y_normalized) where 0,0 is top-left. """ x1, y1, x2, y2 = bbox center_x = (x1 + x2) / 2.0 center_y = (y1 + y2) / 2.0 # Normalize to [0, 1] x_norm = center_x / frame_width if frame_width > 0 else 0.5 y_norm = center_y / frame_height if frame_height > 0 else 0.5 return (x_norm, y_norm) def _find_best_match(self, target_box: Tuple[float, float, float, float], predicted_frame_data: Dict[int, Tuple[Tuple, str]], iou_threshold: float) -> Tuple[Optional[str], float]: """ Find best matching box in predicted frame data using IoU. """ best_iou = 0.0 best_team_id = None for idx, (bbox, team_cls_id) in predicted_frame_data.items(): iou = self._calculate_iou(target_box, bbox) if iou > best_iou and iou >= iou_threshold: best_iou = iou best_team_id = team_cls_id return (best_team_id, best_iou) def _detect_objects_batch(self, decoded_images: List[ndarray]) -> Dict[int, List[BoundingBox]]: 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): """ Hybrid team classification combining: 1. Appearance features (OSNet) 2. Color signatures (HSV/LAB) 3. Spatial priors (left/right side of pitch) 4. Temporal tracking (same player = same team) """ start = time.time() # Phase 1: Collect samples and fit appearance-based classifier fit_sample_size = min(self.MAX_SAMPLES_FOR_FIT, len(detection_results) * 10) 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 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 or cls_id != 2: continue crop = frame_image[int(y1):int(y2), int(x1):int(x2)] if crop.size > 0: player_crops_for_fit.append(crop) if self.team_classifier and not self.team_classifier_fitted and len(player_crops_for_fit) >= fit_sample_size: print(f"Fitting TeamClassifier (OSNet) 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) >= self.MIN_SAMPLES_FOR_FIT: 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 print(f"Fitting time: {time.time() - start:.2f}s") # Phase 2: Hybrid classification for all frames start = time.time() bboxes: dict[int, list[BoundingBox]] = {} # Temporal tracking: {track_id: (team_id, confidence, last_frame)} player_tracks: Dict[Tuple, Tuple[int, float, int]] = {} # Spatial team assignment: track which team is on which side left_side_team = None right_side_team = None for frame_id in range(len(detection_results)): detection_box = detection_results[frame_id].boxes.data frame_image = decoded_images[frame_id] frame_h, frame_w = frame_image.shape[:2] boxes = [] # Collect all players in this frame player_data = [] # (idx, crop, bbox, spatial_pos, color_sig) for idx, box in enumerate(detection_box): x1, y1, x2, y2, conf, cls_id = box.tolist() if cls_id != 2 or conf < 0.6: continue crop = frame_image[int(y1):int(y2), int(x1):int(x2)] if crop.size == 0: continue bbox = (x1, y1, x2, y2) spatial_pos = self._get_spatial_position(bbox, frame_w, frame_h) color_sig = self._extract_color_signature(crop) player_data.append((idx, crop, bbox, spatial_pos, color_sig)) if len(player_data) == 0: bboxes[offset + frame_id] = [] continue # Step 1: Get appearance-based predictions (OSNet) appearance_predictions = {} if self.team_classifier and self.team_classifier_fitted: crops = [data[1] for data in player_data] appearance_team_ids = self.team_classifier.predict(crops) for (idx, _, _, _, _), team_id in zip(player_data, appearance_team_ids): appearance_predictions[idx] = team_id # Step 2: Extract color signatures and cluster color_signatures = [] color_indices = [] for idx, _, _, _, color_sig in player_data: if color_sig is not None: color_signatures.append(color_sig) color_indices.append(idx) color_predictions = {} if len(color_signatures) >= 4: try: from sklearn.cluster import KMeans color_kmeans = KMeans(n_clusters=2, random_state=42, n_init=10) color_clusters = color_kmeans.fit_predict(color_signatures) for idx, cluster_id in zip(color_indices, color_clusters): color_predictions[idx] = cluster_id except Exception as e: print(f"Color clustering failed: {e}") # Step 3: Apply spatial priors # Determine which team is on which side based on majority if left_side_team is None or right_side_team is None: left_side_players = [p for p in player_data if p[3][0] < 0.5] # x < 0.5 right_side_players = [p for p in player_data if p[3][0] >= 0.5] # x >= 0.5 if len(left_side_players) >= 2 and len(right_side_players) >= 2: # Use appearance predictions to determine sides left_teams = [appearance_predictions.get(p[0]) for p in left_side_players if p[0] in appearance_predictions] right_teams = [appearance_predictions.get(p[0]) for p in right_side_players if p[0] in appearance_predictions] if left_teams and right_teams: left_team_mode = max(set(left_teams), key=left_teams.count) right_team_mode = max(set(right_teams), key=right_teams.count) if left_team_mode != right_team_mode: left_side_team = left_team_mode right_side_team = right_team_mode # Step 4: Combine predictions with voting final_predictions = {} for idx, _, bbox, spatial_pos, _ in player_data: votes = [] weights = [] # Appearance vote (weight: 0.4) if idx in appearance_predictions: votes.append(appearance_predictions[idx]) weights.append(0.4) # Color vote (weight: 0.3) if idx in color_predictions: votes.append(color_predictions[idx]) weights.append(0.3) # Spatial vote (weight: 0.3) if left_side_team is not None and right_side_team is not None: x_pos, _ = spatial_pos if x_pos < 0.5: spatial_team = left_side_team else: spatial_team = right_side_team votes.append(spatial_team) weights.append(0.3) # Temporal vote (weight: 0.2) - match with previous frames if len(votes) > 0: # Simple temporal matching: find similar bbox in previous frames best_track_match = None best_track_iou = 0.0 for track_key, (track_team, track_conf, track_frame) in player_tracks.items(): if abs(track_frame - frame_id) <= 5: # Within 5 frames track_bbox = track_key iou = self._calculate_iou(bbox, track_bbox) if iou > best_track_iou and iou > 0.3: best_track_iou = iou best_track_match = track_team if best_track_match is not None: votes.append(best_track_match) weights.append(0.2) # Weighted voting if len(votes) > 0: team_0_score = sum(w for v, w in zip(votes, weights) if v == 0) team_1_score = sum(w for v, w in zip(votes, weights) if v == 1) if team_0_score > team_1_score: final_team = 0 elif team_1_score > team_0_score: final_team = 1 else: # Tie: use appearance prediction or first vote final_team = votes[0] if votes else 0 final_predictions[idx] = final_team # Update tracking track_key = bbox player_tracks[track_key] = (final_team, max(team_0_score, team_1_score), frame_id) # Step 5: Generate output boxes for idx, box in enumerate(detection_box): x1, y1, x2, y2, conf, cls_id = box.tolist() if cls_id == 2 and conf < 0.6: continue # Check overlap with staff 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 with team prediction if idx in final_predictions: mapped_cls_id = str(6 + int(final_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 print(f"Hybrid team classification time: {time.time() - start:.2f}s") 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}") # Use hybrid team classification start = time.time() bboxes = self._team_classify(detection_results, batch_images, offset) end = time.time() print(f"Team classify time: {end - start}") # Phase 3: Keypoint Detection # keypoints: Dict[int, List[Tuple[int, int]]] = {} # keypoints = self._detect_keypoints_batch(batch_images, offset, n_keypoints) 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, []) result = TVFrameResult( frame_id=frame_number, boxes=frame_boxes, keypoints=keypoints.get( frame_number, [(0, 0) for _ in range(n_keypoints)], ), ) results.append(result) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() return results def _detect_keypoints_batch(self, batch_images: List[ndarray], offset: int, n_keypoints: int) -> Dict[int, List[Tuple[int, int]]]: """ Phase 3: Keypoint detection for all frames in batch. Args: batch_images: List of images to process offset: Frame offset for numbering n_keypoints: Number of keypoints expected Returns: Dictionary mapping frame_id to list of keypoint coordinates """ keypoints: Dict[int, List[Tuple[int, int]]] = {} keypoints_model_results = self.keypoints_model.predict(batch_images) if keypoints_model_results is None: return keypoints for frame_idx_in_batch, detection in enumerate(keypoints_model_results): if not hasattr(detection, "keypoints") or detection.keypoints is None: continue # Extract keypoints with confidence 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)) # Pad or truncate to expected number of keypoints 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] # Filter keypoints based on confidence thresholds filtered_keypoints: List[Tuple[int, int]] = [] for idx, (x, y, confidence) in enumerate(frame_keypoints_with_conf): if idx in self.CORNER_INDICES: # Corner keypoints have lower confidence threshold if confidence < 0.3: filtered_keypoints.append((0, 0)) else: filtered_keypoints.append((int(x), int(y))) else: # Regular keypoints if confidence < 0.5: filtered_keypoints.append((0, 0)) else: filtered_keypoints.append((int(x), int(y))) frame_id = offset + frame_idx_in_batch keypoints[frame_id] = filtered_keypoints return keypoints