miner.py

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:
    """
    Enhanced HSV-based team classifier with temporal consistency and confidence weighting.
    Fast and lightweight, suitable for real-time processing.
    """
    
    def __init__(self, hue_pivot: float = 90.0, temporal_weight: float = 0.3):
        """
        Initialize HSV-based team classifier.
        
        Args:
            hue_pivot: Hue threshold for single player classification (default: 90.0)
            temporal_weight: Weight for temporal consistency (0.0-1.0)
        """
        self.hue_pivot = hue_pivot
        self.temporal_weight = temporal_weight
        self.cluster_centers: np.ndarray | None = None
        self.previous_assignments: Dict[int, int] = {}  # bbox_id -> team_id
        self.assignment_confidence: Dict[int, float] = {}  # bbox_id -> confidence
    
    @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 _calculate_bbox_similarity(self, box1, box2) -> float:
        """Calculate similarity between two bounding boxes based on center distance."""
        center1 = ((box1.x1 + box1.x2) / 2, (box1.y1 + box1.y2) / 2)
        center2 = ((box2.x1 + box2.x2) / 2, (box2.y1 + box2.y2) / 2)
        distance = np.sqrt((center1[0] - center2[0])**2 + (center1[1] - center2[1])**2)
        # Normalize by image diagonal (assuming 1920x1080)
        max_distance = np.sqrt(1920**2 + 1080**2)
        return max(0, 1 - distance / max_distance)
    
    def _apply_temporal_consistency(
        self, 
        current_labels: np.ndarray, 
        boxes: List,
        hsv_features: np.ndarray
    ) -> np.ndarray:
        """Apply temporal consistency to reduce team assignment flickering."""
        if not self.previous_assignments:
            return current_labels
        
        adjusted_labels = current_labels.copy()
        
        for i, (box, current_label) in enumerate(zip(boxes, current_labels)):
            best_match_id = None
            best_similarity = 0.0
            
            # Find best matching previous bbox
            for prev_id, prev_team in self.previous_assignments.items():
                # In real implementation, you'd track bbox IDs across frames
                # For now, use position-based matching
                similarity = 0.8  # Placeholder - would use actual bbox tracking
                if similarity > best_similarity and similarity > 0.5:
                    best_similarity = similarity
                    best_match_id = prev_id
            
            # Apply temporal consistency if confident match found
            if best_match_id and best_similarity > 0.7:
                prev_confidence = self.assignment_confidence.get(best_match_id, 0.5)
                current_confidence = 0.8  # Based on HSV feature quality
                
                if prev_confidence > current_confidence * 1.2:
                    adjusted_labels[i] = self.previous_assignments[best_match_id]
        
        return adjusted_labels

    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 with temporal consistency.
        
        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)
            
            # Apply temporal consistency
            if self.temporal_weight > 0:
                labels = self._apply_temporal_consistency(labels, boxes, hsv_features)
            
            # Update tracking
            for i, (box, label) in enumerate(zip(boxes, labels)):
                bbox_id = hash((box.x1, box.y1, box.x2, box.y2)) % 10000  # Simple ID
                self.previous_assignments[bbox_id] = int(label)
                self.assignment_confidence[bbox_id] = 0.8
            
            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)

    def _multi_scale_detection(self, img_bgr: np.ndarray) -> List[BoundingBox]:
        """
        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 = [1.0, 1.15, 0.85]  # Conservative scales for better stability
        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))
            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, yolo_cls_id = box.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
                    
                    # 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)
                    
                    # Boost confidence for detections at optimal scales
                    box_area = (x2 - x1) * (y2 - y1)
                    if scale == 1.15 and box_area < 2500:  # Small objects benefit from upscaling
                        conf *= 1.08
                    elif scale == 0.85 and box_area > 8000:  # Large objects benefit from downscaling
                        conf *= 1.03
                    
                    all_detections.append(BoundingBox(
                        x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2),
                        cls_id=validator_cls_id, conf=float(conf)
                    ))
        
        # Apply multi-scale NMS
        return self._multi_scale_nms(all_detections)
    
    def _multi_scale_nms(self, boxes: List[BoundingBox], iou_threshold: float = 0.45) -> List[BoundingBox]:
        """
        Multi-scale Non-Maximum Suppression that preserves detections from different scales.
        """
        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
                remaining = []
                for box in class_boxes_sorted:
                    iou = self._calculate_iou(current, box)
                    if iou < iou_threshold:
                        remaining.append(box)
                    elif box.conf > current.conf * 0.92:  # Keep if confidence is very close
                        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

    # ========================================================================
    # 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 with multi-scale approach
        if process_objects:
            for frame_idx, frame_image in enumerate(batch_images):
                # Use multi-scale detection for better small object detection
                boxes = self._multi_scale_detection(frame_image)
                
                # 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