miner.py

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