from pathlib import Path from typing import Generator, Iterable, List, TypeVar, Tuple, Dict, Literal, Optional from ultralytics import YOLO from numpy import ndarray from pydantic import BaseModel import numpy as np import torch import torchvision.transforms as T from sklearn.cluster import KMeans import torchvision.models as models import cv2 # ============================================================================ # Utility Functions for Batching # ============================================================================ V = TypeVar("V") def create_batches( sequence: Iterable[V], batch_size: int ) -> Generator[List[V], None, None]: """ Generate batches from a sequence with a specified batch size. Args: sequence (Iterable[V]): The input sequence to be batched. batch_size (int): The size of each batch. Yields: Generator[List[V], None, None]: A generator yielding batches of the input sequence. """ batch_size = max(batch_size, 1) current_batch = [] for element in sequence: if len(current_batch) == batch_size: yield current_batch current_batch = [] current_batch.append(element) if current_batch: yield current_batch # ============================================================================ # Team Classification using HSV Color Space # ============================================================================ class HSVTeamClassifier: """ A classifier that uses HSV color space features for team classification. Fast and lightweight, suitable for real-time processing. """ def __init__(self, hue_pivot: float = 90.0): """ Initialize HSV-based team classifier. Args: hue_pivot: Hue threshold for single player classification (default: 90.0) """ self.hue_pivot = hue_pivot self.cluster_centers: np.ndarray | None = None @staticmethod def _extract_hsv_features_from_crop(img_bgr: np.ndarray) -> Tuple[float, float]: """ Extract mean hue and saturation from an image crop. Args: img_bgr: BGR image crop Returns: Tuple of (mean_hue, mean_saturation) """ if img_bgr.size == 0: return (0.0, 0.0) hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV) mean_hue = float(np.mean(hsv[:, :, 0])) mean_saturation = float(np.mean(hsv[:, :, 1])) return (mean_hue, mean_saturation) def _extract_hsv_features_with_green_filter( self, img_bgr: np.ndarray, box, img_width: int, img_height: int ) -> np.ndarray: """ Extract HSV features from ROI, filtering out green (grass) pixels. Args: img_bgr: Full frame image box: Bounding box to extract ROI from img_width, img_height: Image dimensions Returns: Array of [hue, saturation] features """ x1, y1, x2, y2 = Miner._clip_box_to_image( box.x1, box.y1, box.x2, box.y2, img_width, img_height ) 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) # Filter out green pixels (grass) 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_pixels = hsv.shape[0] * hsv.shape[1] # If enough non-green pixels, use only those if num_non_green > max(50, total_pixels // 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: # Use all pixels if too few non-green h_mean, s_mean = self._extract_hsv_features_from_crop(roi) return np.array([h_mean, s_mean], dtype=np.float32) def _cluster_players_hsv( self, hsv_features: np.ndarray ) -> Tuple[np.ndarray, np.ndarray]: """ Cluster players into two teams using K-means on HSV features. Args: hsv_features: Array of HSV features (N, 2) Returns: Tuple of (labels, cluster_centers) """ if len(hsv_features) < 2: return np.array([]), np.array([]) criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 20, 1.0) _, labels, centers = cv2.kmeans( np.float32(hsv_features), K=2, bestLabels=None, criteria=criteria, attempts=5, flags=cv2.KMEANS_PP_CENTERS, ) # Sort centers by hue for consistent team assignment order = np.argsort(centers[:, 0]) centers_sorted = centers[order] remap = {old_idx: new_idx for new_idx, old_idx in enumerate(order)} labels_remapped = np.vectorize(remap.get)(labels.reshape(-1)) return labels_remapped, centers_sorted def predict( self, crops: List[np.ndarray], boxes: List, frame_image: ndarray ) -> Tuple[np.ndarray, np.ndarray | None]: """ Predict team labels for player crops using HSV features. Args: crops: List of player image crops boxes: List of corresponding bounding boxes frame_image: Full frame image for feature extraction Returns: Tuple of (team_labels, cluster_centers) """ if len(crops) == 0: return np.array([]), None h, w = frame_image.shape[:2] hsv_features = [] for box in boxes: features = self._extract_hsv_features_with_green_filter( frame_image, box, w, h ) hsv_features.append(features) hsv_features = np.vstack(hsv_features) if len(hsv_features) >= 2: labels, centers = self._cluster_players_hsv(hsv_features) self.cluster_centers = centers return labels, centers elif len(hsv_features) == 1: # Single player: use hue pivot hue = hsv_features[0, 0] label = 0 if float(hue) < self.hue_pivot else 1 return np.array([label]), None else: return np.array([]), None # ============================================================================ # Team Classification using ResNet18 Features # ============================================================================ class ResNetTeamClassifier: """ A classifier that uses ResNet18 for feature extraction and KMeans for clustering. """ def __init__(self, device: str = 'cpu', batch_size: int = 32): """ Initialize the TeamClassifier with device and batch size. Args: device (str): The device to run the model on ('cpu' or 'cuda'). batch_size (int): The batch size for processing images. """ self.device = device self.batch_size = batch_size # Load pretrained ResNet18 (fix deprecation warning) self.features_model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) # Remove the final classification layer to get features self.features_model = torch.nn.Sequential(*list(self.features_model.children())[:-1]) self.features_model.to(device) self.features_model.eval() # Create transform using torchvision (works directly with numpy arrays) self.transform = T.Compose([ T.ToTensor(), # numpy (H,W,C) BGR → tensor (C,H,W) RGB, scale to [0,1] T.Resize((224, 224)), T.Normalize( mean=[0.485, 0.456, 0.406], # ImageNet pretrained means std=[0.229, 0.224, 0.225] # ImageNet pretrained stds ) ]) # KMeans clustering self.cluster_model = KMeans(n_clusters=2, random_state=42) def extract_features(self, crops: List[np.ndarray]) -> np.ndarray: """ Extract features from a list of image crops using ResNet18. Args: crops (List[np.ndarray]): List of image crops (CV2 numpy arrays, BGR format). Returns: np.ndarray: Extracted features as a numpy array (N, 512). """ # Batch numpy arrays directly (no PIL conversion) batches = create_batches(crops, self.batch_size) embeddings = [] with torch.no_grad(): for batch in batches: # Transform numpy arrays directly to tensors and stack inputs = torch.stack([self.transform(crop) for crop in batch]).to(self.device) # Extract features from ResNet18 (before final classification layer) features = self.features_model(inputs) # Flatten the features (batch_size, 512, 1, 1) -> (batch_size, 512) features = features.view(features.size(0), -1) embeddings.append(features.cpu().numpy()) return np.concatenate(embeddings) def fit(self, crops: List[np.ndarray], max_samples: int = 100) -> None: """ Fit the classifier model on a list of image crops. Args: crops (List[np.ndarray]): List of image crops. max_samples (int): Maximum number of samples to use for fitting. """ # Random sample if too many crops if len(crops) > max_samples: indices = np.random.choice(len(crops), max_samples, replace=False) crops = [crops[i] for i in indices] # Extract features (512D embeddings from ResNet18) embeddings = self.extract_features(crops) # Fit KMeans directly (no UMAP) self.cluster_model.fit(embeddings) def predict(self, crops: List[np.ndarray]) -> np.ndarray: """ Predict the cluster labels for a list of image crops. Args: crops (List[np.ndarray]): List of image crops. Returns: np.ndarray: Predicted cluster labels (0 or 1). """ if len(crops) == 0: return np.array([]) # Extract features (512D embeddings from ResNet18) embeddings = self.extract_features(crops) # Predict directly (no UMAP transform) return self.cluster_model.predict(embeddings) # ============================================================================ # Data Models # ============================================================================ 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]] # ============================================================================ # Main Miner Class # ============================================================================ class Miner: """ Enhanced miner combining best practices from v1 and competitor's v3. Features: - Multiple team classification methods (HSV, ResNet, ensemble) - Two-stage box suppression (quasi-total containment + small contained) - Simplified multiple goalkeeper handling (confidence-based) - Proper task_type support for selective processing - Boundary-aware box clipping """ # Constants for box suppression (from competitor's approach) 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, team_classification_method: Literal["hsv", "resnet", "ensemble"] = "hsv" ) -> None: """ Loads all ML models from the repository. Args: path_hf_repo (Path): Path to the downloaded HuggingFace Hub repository team_classification_method (str): Method for team classification - "hsv": Fast HSV-based classification (default) - "resnet": Robust ResNet18-based classification - "ensemble": Combine both methods (vote-based) """ self.bbox_model = YOLO(path_hf_repo / "detection.pt") print(f"✅ BBox Model Loaded") self.keypoints_model = YOLO(path_hf_repo / "keypoint.pt") print(f"✅ Keypoints Model Loaded") # Initialize team classification method self.team_classification_method = team_classification_method if team_classification_method == "hsv": self.hsv_classifier = HSVTeamClassifier(hue_pivot=self.SINGLE_PLAYER_HUE_PIVOT) self.resnet_classifier = None self.team_classifier_fitted = False # HSV doesn't need fitting print(f"✅ HSV Team Classifier Initialized") elif team_classification_method == "resnet": device = "cuda" if torch.cuda.is_available() else "cpu" print(f"🔧 Using device: {device}") self.resnet_classifier = ResNetTeamClassifier(device=device, batch_size=32) self.hsv_classifier = None self.team_classifier_fitted = False print(f"✅ ResNet Team Classifier Loaded") elif team_classification_method == "ensemble": device = "cuda" if torch.cuda.is_available() else "cpu" print(f"🔧 Using device: {device}") self.hsv_classifier = HSVTeamClassifier(hue_pivot=self.SINGLE_PLAYER_HUE_PIVOT) self.resnet_classifier = ResNetTeamClassifier(device=device, batch_size=32) self.team_classifier_fitted = False # Only ResNet needs fitting print(f"✅ Ensemble Team Classifiers Loaded (HSV + ResNet)") else: raise ValueError( f"Invalid team_classification_method: {team_classification_method}. " "Must be 'hsv', 'resnet', or 'ensemble'" ) def __repr__(self) -> str: """Information about miner returned in the health endpoint.""" classifier_info = f"Team Classification: {self.team_classification_method}" if self.team_classification_method == "hsv": classifier_info += f" ({type(self.hsv_classifier).__name__})" elif self.team_classification_method == "resnet": classifier_info += f" ({type(self.resnet_classifier).__name__})" else: classifier_info += f" (HSV + ResNet)" return ( f"BBox Model: {type(self.bbox_model).__name__}\n" f"Keypoints Model: {type(self.keypoints_model).__name__}\n" f"{classifier_info}" ) # ======================================================================== # Post-processing Helper Methods # ======================================================================== @staticmethod def _map_yolo_to_validator_cls_id(yolo_cls_id: int) -> int | None: """ Map YOLO model class ID (new model format) to validator format. YOLO model mapping: 0: 'Player', 1: 'GoalKeeper', 2: 'Ball', 3: 'Main Referee', 4: 'Side Referee', 5: 'Staff Member' Validator format: 0: 'ball', 1: 'goalkeeper', 2: 'player', 3: 'referee', 6: 'team1', 7: 'team2' Args: yolo_cls_id: Class ID from YOLO model Returns: Mapped class ID in validator format, or None if should be skipped """ if yolo_cls_id == 0: # YOLO Player -> Validator Player (2) return 2 elif yolo_cls_id == 1: # YOLO GoalKeeper -> Validator GoalKeeper (1) return 1 elif yolo_cls_id == 2: # YOLO Ball -> Validator Ball (0) return 0 elif yolo_cls_id in [3, 4]: # YOLO Main/Side Referee -> Validator Referee (3) return 3 else: # Staff Member or other -> skip return None @staticmethod def _clip_box_to_image(x1: int, y1: int, x2: int, y2: int, w: int, h: int) -> Tuple[int, int, int, int]: """ Clip bounding box coordinates to ensure they are within image boundaries. (Adopted from competitor's approach - simpler and more efficient) Args: x1, y1, x2, y2: Box coordinates w, h: Image dimensions Returns: Clipped coordinates (x1, y1, x2, y2) """ 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: """Calculate the area of a bounding box.""" return max(0, bb.x2 - bb.x1) * max(0, bb.y2 - bb.y1) @staticmethod def _intersect_area(a: BoundingBox, b: BoundingBox) -> int: """Calculate the intersection area between two boxes.""" 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) def _ioa(self, a: BoundingBox, b: BoundingBox) -> float: """ Calculate Intersection over Area (IoA) of box a in box b. (Adopted from competitor's approach) """ 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]: """ Remove boxes that are almost completely contained within another box. (Adopted from competitor's approach - cleaner separation of concerns) Strategy: If box_i is >= 90% contained in box_j, remove box_i. This handles cases where one box is a near-duplicate of another. """ 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]: """ Remove small boxes that are significantly contained within larger boxes. (Adopted from competitor's approach - cleaner separation of concerns) Strategy: If a small box (<= 50% size) is >= 85% contained in a larger box, remove the small box (likely a duplicate detection). """ 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 _handle_multiple_balls( self, all_boxes: List[BoundingBox] ) -> List[BoundingBox]: """ When multiple footballs are detected, keep only the one with highest confidence. """ ball_detections = [box for box in all_boxes if box.cls_id == 0] if len(ball_detections) <= 1: return all_boxes # Find the ball with highest confidence best_ball = max(ball_detections, key=lambda b: b.conf) # Remove all balls, then add back the best one filtered_boxes = [box for box in all_boxes if box.cls_id != 0] filtered_boxes.append(best_ball) return filtered_boxes def _reclass_extra_goalkeepers( self, img_bgr: np.ndarray, boxes: List[BoundingBox], cluster_centers: Optional[np.ndarray], ) -> None: """ When multiple goalkeepers are detected, keep the one with highest confidence and reclassify the rest as regular players. (Adopted from competitor's simpler approach - confidence-based selection) Args: img_bgr: Current frame image boxes: List of all detected boxes (modified in-place) cluster_centers: Pre-computed team cluster centers (if available) """ gk_idxs = [i for i, bb in enumerate(boxes) if int(bb.cls_id) == 1] if len(gk_idxs) <= 1: return # Sort by confidence and keep the highest 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:] # Reclassify extra goalkeepers for gki in to_reclass: # Extract HSV features for team assignment h, w = img_bgr.shape[:2] hs_gk = self.hsv_classifier._extract_hsv_features_with_green_filter( img_bgr, boxes[gki], w, h ) if self.hsv_classifier else None # Assign team based on available classifier and cluster centers if cluster_centers is not None and len(cluster_centers) >= 2: if self.team_classification_method == "resnet" and self.team_classifier_fitted: # Use ResNet features if available try: x1, y1, x2, y2 = self._clip_box_to_image( boxes[gki].x1, boxes[gki].y1, boxes[gki].x2, boxes[gki].y2, w, h ) gk_crop = img_bgr[y1:y2, x1:x2] if gk_crop.size > 0: gk_features = self.resnet_classifier.extract_features([gk_crop])[0] d0 = float(np.linalg.norm(gk_features - cluster_centers[0])) d1 = float(np.linalg.norm(gk_features - cluster_centers[1])) assign_cls = 6 if d0 <= d1 else 7 else: assign_cls = 6 except Exception: # Fallback to HSV if ResNet fails if hs_gk 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 else: # Use HSV features if hs_gk 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 else: # No cluster centers - use hue pivot or default if hs_gk is not None: assign_cls = 6 if float(hs_gk[0]) < self.SINGLE_PLAYER_HUE_PIVOT else 7 else: assign_cls = 6 boxes[gki].cls_id = int(assign_cls) # ======================================================================== # Main Prediction Methods # ======================================================================== def fit_team_classifier( self, batch_images: list[ndarray], player_class_id: int = 0 ) -> None: """ Fit the team classifier on player crops from batch images. Only needed for ResNet or ensemble methods. Args: batch_images: List of images to extract player crops from player_class_id: YOLO class ID that represents players (default: 0 for new model) """ if self.team_classification_method == "hsv": print("ℹī¸ HSV classifier doesn't require fitting") return player_crops = [] bbox_model_results = self.bbox_model.predict(batch_images) if bbox_model_results is not None: for frame_idx, detection in enumerate(bbox_model_results): if not hasattr(detection, "boxes") or detection.boxes is None: continue frame_image = batch_images[frame_idx] h, w = frame_image.shape[:2] for box in detection.boxes.data: x1, y1, x2, y2, conf, yolo_cls_id = box.tolist() # Only collect crops for players (YOLO class ID 0 = Player) if int(yolo_cls_id) == player_class_id: x1_clip, y1_clip, x2_clip, y2_clip = self._clip_box_to_image( int(x1), int(y1), int(x2), int(y2), w, h ) crop = frame_image[y1_clip:y2_clip, x1_clip:x2_clip] if crop.size > 0: player_crops.append(crop) if len(player_crops) > 0: if self.team_classification_method == "resnet": self.resnet_classifier.fit(player_crops) self.team_classifier_fitted = True print(f"✅ ResNet team classifier fitted on {len(player_crops)} player crops") elif self.team_classification_method == "ensemble": self.resnet_classifier.fit(player_crops) self.team_classifier_fitted = True print(f"✅ ResNet classifier (in ensemble) fitted on {len(player_crops)} player crops") else: print("⚠ī¸ No player crops found to fit team classifier") def predict_batch( self, batch_images: list[ndarray], offset: int, n_keypoints: int, task_type: Optional[str] = None, ) -> list[TVFrameResult]: """ Miner prediction for a batch of images with enhanced post-processing. Args: batch_images (list[np.ndarray]): A list of images to process offset (int): Frame number of the first image in the batch n_keypoints (int): Number of keypoints expected per frame task_type (str | None): - None: Process both object and keypoint detection - "object": Only process object detection - "keypoint": Only process keypoint detection Returns: list[TVFrameResult]: Predictions for each image in the batch """ # Determine which tasks to process (adopted from competitor's approach) 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]] = {} # Process object detection 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, detection in enumerate(bbox_model_results): if not hasattr(detection, "boxes") or detection.boxes is None: continue frame_image = batch_images[frame_idx] h, w = frame_image.shape[:2] detection_boxes = detection.boxes.data # Convert to BoundingBox objects (don't normalize yet - clip when needed) # Map YOLO class IDs to validator format boxes = [] for box_data in detection_boxes: x1, y1, x2, y2, conf, yolo_cls_id = box_data.tolist() # Map YOLO class ID to validator format validator_cls_id = self._map_yolo_to_validator_cls_id(int(yolo_cls_id)) if validator_cls_id is None: continue # Skip Staff Member and other classes boxes.append( BoundingBox( x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2), cls_id=validator_cls_id, conf=float(conf), ) ) # Handle multiple balls first boxes = self._handle_multiple_balls(boxes) # Apply two-stage suppression (adopted from competitor's approach) boxes = self.suppress_quasi_total_containment(boxes) boxes = self.suppress_small_contained(boxes) # Team classification for players player_boxes = [box for idx, box in enumerate(boxes) if box.cls_id == 2] player_indices = [idx for idx, box in enumerate(boxes) if box.cls_id == 2] team_cluster_centers = None team_labels = None if len(player_boxes) > 0: if self.team_classification_method == "hsv": # HSV-based classification (no fitting needed) player_crops = [ frame_image[box.y1:box.y2, box.x1:box.x2] for box in player_boxes ] team_labels, team_cluster_centers = self.hsv_classifier.predict( player_crops, player_boxes, frame_image ) elif self.team_classification_method == "resnet": # ResNet-based classification (requires fitting) if self.team_classifier_fitted: player_crops = [ frame_image[box.y1:box.y2, box.x1:box.x2] for box in player_boxes ] team_labels = self.resnet_classifier.predict(player_crops) # Get cluster centers if hasattr(self.resnet_classifier.cluster_model, 'cluster_centers_'): team_cluster_centers = self.resnet_classifier.cluster_model.cluster_centers_ elif self.team_classification_method == "ensemble": # Ensemble: combine HSV and ResNet predictions player_crops = [ frame_image[box.y1:box.y2, box.x1:box.x2] for box in player_boxes ] # Get predictions from both methods hsv_labels, hsv_centers = self.hsv_classifier.predict( player_crops, player_boxes, frame_image ) resnet_labels = None resnet_centers = None if self.team_classifier_fitted: resnet_labels = self.resnet_classifier.predict(player_crops) if hasattr(self.resnet_classifier.cluster_model, 'cluster_centers_'): resnet_centers = self.resnet_classifier.cluster_model.cluster_centers_ # Combine predictions (vote-based) if resnet_labels is not None and len(resnet_labels) == len(hsv_labels): # Vote: if both agree, use that; otherwise prefer ResNet team_labels = np.array([ resnet_labels[i] if resnet_labels[i] == hsv_labels[i] else resnet_labels[i] # Prefer ResNet on disagreement for i in range(len(hsv_labels)) ]) team_cluster_centers = resnet_centers # Use ResNet centers else: # Fallback to HSV if ResNet not available team_labels = hsv_labels team_cluster_centers = hsv_centers # Update player class IDs to team IDs (6 or 7) if team_labels is not None and len(team_labels) == len(player_indices): for idx, team_label in zip(player_indices, team_labels): boxes[idx].cls_id = 6 + int(team_label) # Handle multiple goalkeepers (simplified approach from competitor) self._reclass_extra_goalkeepers( frame_image, boxes, team_cluster_centers ) bboxes[offset + frame_idx] = boxes # Process keypoint detection 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, detection in enumerate(keypoints_model_results): if not hasattr(detection, "keypoints") or detection.keypoints is None: continue frame_keypoints: list[tuple[int, int, float]] = [] for i, part_points in enumerate(detection.keypoints.data): for k_id, (x, y, _) in enumerate(part_points): confidence = detection.keypoints.conf[i][k_id] frame_keypoints.append((int(x), int(y), float(confidence))) # Pad or truncate to match expected number of keypoints if len(frame_keypoints) < n_keypoints: frame_keypoints.extend( [(0, 0, 0.0)] * (n_keypoints - len(frame_keypoints)) ) else: frame_keypoints = frame_keypoints[:n_keypoints] # Filter keypoints based on confidence # Corner keypoints use lower threshold (0.3) to ensure homography can be computed filtered_keypoints = [] for idx, (x, y, confidence) in enumerate(frame_keypoints): if idx in self.CORNER_INDICES: # For corner keypoints, use lower threshold if confidence < 0.3: filtered_keypoints.append((0, 0)) else: filtered_keypoints.append((int(x), int(y))) else: # For non-corner keypoints, use standard threshold if confidence < 0.5: filtered_keypoints.append((0, 0)) else: filtered_keypoints.append((int(x), int(y))) keypoints[offset + frame_idx] = filtered_keypoints # Combine results 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