introduce osnet

config.yml

@@ -2,13 +2,15 @@ Image:
   from_base: parachutes/python:3.12
   run_command:
     - pip install --upgrade setuptools wheel
+    - pip install "torch==2.7.1" "torchvision==0.22.1"
     - pip install "ultralytics==8.3.222" "opencv-python-headless" "numpy" "pydantic"
-    - pip install "tensorflow" "torch==2.7.1" "torchvision==0.22.1" "torch-tensorrt==2.7"
+    - pip install scikit-learn
+    - pip install onnxruntime-gpu
   set_workdir: /app
 
 NodeSelector:
   gpu_count: 1
-  min_vram_gb_per_gpu: 16
+  min_vram_gb_per_gpu: 24
   exclude:
     - "5090"
     - b200
@@ -19,4 +21,5 @@ Chute:
   timeout_seconds: 900
   concurrency: 4
   max_instances: 5
-  scaling_threshold: 0.5
+  scaling_threshold: 0.5
+  shutdown_after_seconds: 3600

miner.py

@@ -4,41 +4,22 @@ import sys
 import os
 
 from numpy import ndarray
-import numpy as np
 from pydantic import BaseModel
-import cv2
-
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
-os.environ["OMP_NUM_THREADS"] = "16"
-os.environ["TF_NUM_INTRAOP_THREADS"] = "16"
-os.environ["TF_NUM_INTEROP_THREADS"] = "2"
-os.environ["CUDA_LAUNCH_BLOCKING"] = "0"
-os.environ["ORT_LOGGING_LEVEL"] = "3"
-os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
-
-import logging
-import tensorflow as tf
-from tensorflow.keras import mixed_precision
-import torch._dynamo
+from ultralytics import YOLO
+from team_cluster import TeamClassifier
+from utils import (
+    BoundingBox, 
+    Constants,
+)
+
+import time
 import torch
-# import torch_tensorrt
 import gc
-from ultralytics import YOLO
 from pitch import process_batch_input, get_cls_net
 import yaml
 
-logging.getLogger("tensorflow").setLevel(logging.ERROR)
-tf.config.threading.set_intra_op_parallelism_threads(16)
-tf.config.threading.set_inter_op_parallelism_threads(2)
-tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
-tf.get_logger().setLevel("ERROR")
-tf.autograph.set_verbosity(0)
-mixed_precision.set_global_policy("mixed_float16")
-tf.config.optimizer.set_jit(True)
-torch._dynamo.config.suppress_errors = True
-
 
 class BoundingBox(BaseModel):
     x1: int
@@ -56,260 +37,340 @@ class TVFrameResult(BaseModel):
 
 
 class Miner:
-    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
-
-    CLS_MAP = {
-        0: 0,
-        1: 1,
-        2: 6,
-        3: 7,
-        4: 3,
-    }
+    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 = 600  # Maximum samples to avoid overfitting
 
     def __init__(self, path_hf_repo: Path) -> None:
-        print(path_hf_repo / "best.pt")
-        self.bbox_model = YOLO(path_hf_repo / "best.pt")
-        print(" BBox Model (best.pt) Loaded")
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        # model_kp_path = path_hf_repo / "SV_kp.engine"
-        # model_kp = torch_tensorrt.load(model_kp_path)
-
-        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()
-
-        # @torch.inference_mode()
-        # def run_inference(model, input_tensor: torch.Tensor):
-        #     input_tensor = input_tensor.to(device).to(memory_format=torch.channels_last)
-        #     output = model.module().forward(input_tensor)
-        #     return output
-
-        # run_inference(model_kp, torch.randn(8, 3, 540, 960, device=device, dtype=torch.float32))
-        self.keypoints_model = model
-        self.kp_threshold = 0.1
-        self.pitch_batch_size = 8
-        print("✅ Keypoints Model Loaded")
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            model_path = path_hf_repo / "football_object_detection.onnx"
+            self.bbox_model = YOLO(model_path)
+            
+            print("BBox Model Loaded")
+
+            team_model_path = path_hf_repo / "osnet_model.pth.tar-100"
+            self.team_classifier = TeamClassifier(
+                device=device,
+                batch_size=32,
+                model_name=str(team_model_path)
+            )
+            print("Team Classifier Loaded")
+            
+            # Team classification state
+            self.team_classifier_fitted = False
+            self.player_crops_for_fit = [] 
+
+            model_kp_path = path_hf_repo / 'keypoint'
+            config_kp_path = path_hf_repo / 'hrnetv2_w48.yaml'
+            cfg_kp = yaml.safe_load(open(config_kp_path, 'r'))
+            
+            loaded_state_kp = torch.load(model_kp_path, map_location=device)
+            model = get_cls_net(cfg_kp)
+            model.load_state_dict(loaded_state_kp)
+            model.to(device)
+            model.eval()
+
+            self.keypoints_model = model
+            self.kp_threshold = 0.1
+            self.pitch_batch_size = 4
+            self.health = "healthy"
+            print("✅ Keypoints Model Loaded")
+        except Exception as e:
+            self.health = "❌ Miner initialization failed: " + str(e)
+            print(self.health)
 
     def __repr__(self) -> str:
-        return (
-            f"BBox Model: {type(self.bbox_model).__name__}\n"
-            f"Keypoints Model: {type(self.keypoints_model).__name__}"
-        )
-
-    @staticmethod
-    def _clip_box_to_image(x1: int, y1: int, x2: int, y2: int, w: int, h: int) -> Tuple[int, int, int, int]:
-        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:
-        return max(0, bb.x2 - bb.x1) * max(0, bb.y2 - bb.y1)
-
-    @staticmethod
-    def _intersect_area(a: BoundingBox, b: BoundingBox) -> int:
-        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)
-
-    @staticmethod
-    def _center(bb: BoundingBox) -> Tuple[float, float]:
-        return (0.5 * (bb.x1 + bb.x2), 0.5 * (bb.y1 + bb.y2))
-
-    @staticmethod
-    def _mean_hs(img_bgr: np.ndarray) -> Tuple[float, float]:
-        hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
-        return float(np.mean(hsv[:, :, 0])), float(np.mean(hsv[:, :, 1]))
-
-    def _hs_feature_from_roi(self, img_bgr: np.ndarray, box: BoundingBox) -> np.ndarray:
-        H, W = img_bgr.shape[:2]
-        x1, y1, x2, y2 = self._clip_box_to_image(box.x1, box.y1, box.x2, box.y2, W, H)
-        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)
-        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 = hsv.shape[0] * hsv.shape[1]
-        if num_non_green > max(50, total // 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
+        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:
-            h_mean, s_mean = self._mean_hs(roi)
-        return np.array([h_mean, s_mean], dtype=np.float32)
+            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)
 
-    def _ioa(self, a: BoundingBox, b: BoundingBox) -> float:
-        inter = self._intersect_area(a, b)
-        aa = self._area(a)
-        if aa <= 0:
+        # 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 inter / aa
-
-    def suppress_quasi_total_containment(self, boxes: List[BoundingBox]) -> List[BoundingBox]:
-        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]:
-        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]:
+
+        return intersection_area / union_area
+
+    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):
+        self.team_classifier_fitted = False
+        start = time.time()
+        # Collect player crops from first batch for fitting
+        fit_sample_size = 600
+        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
-            for j in range(len(boxes)):
-                if i == j or not keep[j]:
+            # Collect player boxes for team classification fitting (first batch only)
+            if len(player_crops_for_fit) < fit_sample_size:
+                frame_image = decoded_images[frame_id]
+                for box in detection_box:
+                    x1, y1, x2, y2, conf, cls_id = box.tolist()
+                    if conf < 0.5:
+                        continue
+                    mapped_cls_id = str(int(cls_id))
+                    # Only collect player crops (cls_id = 2)
+                    if mapped_cls_id == '2':
+                        crop = frame_image[int(y1):int(y2), int(x1):int(x2)]
+                        if crop.size > 0:
+                            player_crops_for_fit.append(crop)
+
+            # Fit team classifier after collecting samples
+            if self.team_classifier and not self.team_classifier_fitted and len(player_crops_for_fit) >= fit_sample_size:
+                print(f"Fitting TeamClassifier with {len(player_crops_for_fit)} player crops")
+                self.team_classifier.fit(player_crops_for_fit)
+                self.team_classifier_fitted = True
+                break
+        if not self.team_classifier_fitted and len(player_crops_for_fit) >= 16:
+            print(f"Fallback: Fitting TeamClassifier with {len(player_crops_for_fit)} player crops")
+            self.team_classifier.fit(player_crops_for_fit)
+            self.team_classifier_fitted = True
+        end = time.time()
+        print(f"Fitting Kmeans time: {end - start}")
+
+        # Second pass: predict teams with configurable frame skipping optimization
+        start = time.time()
+
+        # Get configuration for frame skipping
+        prediction_interval = 1  # Default: predict every 2 frames
+        iou_threshold = 0.3
+
+        print(f"Team classification - prediction_interval: {prediction_interval}, iou_threshold: {iou_threshold}")
+
+        # Storage for predicted frame results: {frame_id: {box_idx: (bbox, team_id)}}
+        predicted_frame_data = {}
+
+        # Step 1: Predict for frames at prediction_interval only
+        frames_to_predict = []
+        for frame_id in range(len(detection_results)):
+            if frame_id % prediction_interval == 0:
+                frames_to_predict.append(frame_id)
+
+        print(f"Predicting teams for {len(frames_to_predict)}/{len(detection_results)} frames "
+                    f"(saving {100 - (len(frames_to_predict) * 100 // len(detection_results))}% compute)")
+
+        for frame_id in frames_to_predict:
+            detection_box = detection_results[frame_id].boxes.data
+            frame_image = decoded_images[frame_id]
+
+            # Collect player crops for this frame
+            frame_player_crops = []
+            frame_player_indices = []
+            frame_player_boxes = []
+
+            for idx, box in enumerate(detection_box):
+                x1, y1, x2, y2, conf, cls_id = box.tolist()
+                if cls_id == 2 and conf < 0.6:
                     continue
-                ai, aj = areas[i], areas[j]
-                if ai == 0 or aj == 0:
+                mapped_cls_id = str(int(cls_id))
+
+                # Collect player crops for prediction
+                if self.team_classifier and self.team_classifier_fitted and mapped_cls_id == '2':
+                    crop = frame_image[int(y1):int(y2), int(x1):int(x2)]
+                    if crop.size > 0:
+                        frame_player_crops.append(crop)
+                        frame_player_indices.append(idx)
+                        frame_player_boxes.append((x1, y1, x2, y2))
+
+            # Predict teams for all players in this frame
+            if len(frame_player_crops) > 0:
+                team_ids = self.team_classifier.predict(frame_player_crops)
+                predicted_frame_data[frame_id] = {}
+                for idx, bbox, team_id in zip(frame_player_indices, frame_player_boxes, team_ids):
+                    # Map team_id (0,1) to cls_id (6,7)
+                    team_cls_id = str(6 + int(team_id))
+                    predicted_frame_data[frame_id][idx] = (bbox, team_cls_id)
+
+        # Step 2: Process all frames (interpolate skipped frames)
+        fallback_count = 0
+        interpolated_count = 0
+        bboxes: dict[int, list[BoundingBox]] = {}
+        for frame_id in range(len(detection_results)):
+            detection_box = detection_results[frame_id].boxes.data
+            frame_image = decoded_images[frame_id]
+            boxes = []
+
+            team_predictions = {}
+
+            if frame_id % prediction_interval == 0:
+                # Predicted frame: use pre-computed predictions
+                if frame_id in predicted_frame_data:
+                    for idx, (bbox, team_cls_id) in predicted_frame_data[frame_id].items():
+                        team_predictions[idx] = team_cls_id
+            else:
+                # Skipped frame: interpolate from neighboring predicted frames
+                # Find nearest predicted frames
+                prev_predicted_frame = (frame_id // prediction_interval) * prediction_interval
+                next_predicted_frame = prev_predicted_frame + prediction_interval
+
+                # Collect current frame player boxes
+                for idx, box in enumerate(detection_box):
+                    x1, y1, x2, y2, conf, cls_id = box.tolist()
+                    if cls_id == 2 and conf < 0.6:
+                        continue
+                    mapped_cls_id = str(int(cls_id))
+
+                    if self.team_classifier and self.team_classifier_fitted and mapped_cls_id == '2':
+                        target_box = (x1, y1, x2, y2)
+
+                        # Try to match with previous predicted frame
+                        best_team_id = None
+                        best_iou = 0.0
+
+                        if prev_predicted_frame in predicted_frame_data:
+                            team_id, iou = self._find_best_match(
+                                target_box,
+                                predicted_frame_data[prev_predicted_frame],
+                                iou_threshold
+                            )
+                            if team_id is not None:
+                                best_team_id = team_id
+                                best_iou = iou
+
+                        # Try to match with next predicted frame if available and no good match yet
+                        if best_team_id is None and next_predicted_frame < len(detection_results):
+                            if next_predicted_frame in predicted_frame_data:
+                                team_id, iou = self._find_best_match(
+                                    target_box,
+                                    predicted_frame_data[next_predicted_frame],
+                                    iou_threshold
+                                )
+                                if team_id is not None and iou > best_iou:
+                                    best_team_id = team_id
+                                    best_iou = iou
+
+                        # Track interpolation success
+                        if best_team_id is not None:
+                            interpolated_count += 1
+                        else:
+                            # Fallback: if no match found, predict individually
+                            crop = frame_image[int(y1):int(y2), int(x1):int(x2)]
+                            if crop.size > 0:
+                                team_id = self.team_classifier.predict([crop])[0]
+                                best_team_id = str(6 + int(team_id))
+                                fallback_count += 1
+
+                        if best_team_id is not None:
+                            team_predictions[idx] = best_team_id
+
+            # Parse boxes with team classification
+            for idx, box in enumerate(detection_box):
+                x1, y1, x2, y2, conf, cls_id = box.tolist()
+                if cls_id == 2 and conf < 0.6:
                     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
+
+                # Check overlap with staff box
+                overlap_staff = False
+                for idy, boxy in enumerate(detection_box):
+                    s_x1, s_y1, s_x2, s_y2, s_conf, s_cls_id = boxy.tolist()
+                    if cls_id == 2 and s_cls_id == 4:
+                        staff_iou = self._calculate_iou(box[:4], boxy[:4])
+                        if staff_iou >= 0.8:
+                            overlap_staff = True
                             break
-                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 _assign_players_two_clusters(self, features: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
-        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 20, 1.0)
-        _, labels, centers = cv2.kmeans(
-            np.float32(features),
-            K=2,
-            bestLabels=None,
-            criteria=criteria,
-            attempts=5,
-            flags=cv2.KMEANS_PP_CENTERS,
-        )
-        return labels.reshape(-1), centers
-
-    def _reclass_extra_goalkeepers(self, img_bgr: np.ndarray, boxes: List[BoundingBox], cluster_centers: np.ndarray | None) -> None:
-        gk_idxs = [i for i, bb in enumerate(boxes) if int(bb.cls_id) == 1]
-        if len(gk_idxs) <= 1:
-            return
-        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:]
-        for gki in to_reclass:
-            hs_gk = self._hs_feature_from_roi(img_bgr, boxes[gki])
-            if cluster_centers 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 if float(hs_gk[0]) < self.SINGLE_PLAYER_HUE_PIVOT else 7
-            boxes[gki].cls_id = int(assign_cls)
-
-    def predict_batch(self, batch_images: List[ndarray], offset: int, n_keypoints: int) -> List[TVFrameResult]:
-        bboxes: Dict[int, List[BoundingBox]] = {}
-        bbox_model_results = self.bbox_model.predict(batch_images)
-        if bbox_model_results is not None:
-            for frame_idx_in_batch, detection in enumerate(bbox_model_results):
-                if not hasattr(detection, "boxes") or detection.boxes is None:
+                if overlap_staff:
                     continue
-                boxes: List[BoundingBox] = []
-                for box in detection.boxes.data:
-                    x1, y1, x2, y2, conf, cls_id = box.tolist()
-                    # if cls_id == 3:
-                    #     cls_id = 2
-                    # elif cls_id == 2:
-                    #     cls_id = 3
+
+                mapped_cls_id = str(int(cls_id))
+
+                # Override cls_id for players with team prediction
+                if idx in team_predictions:
+                    mapped_cls_id = team_predictions[idx]
+                if mapped_cls_id != '4':
+                    if int(mapped_cls_id) == 3 and conf < 0.5:
+                        continue
                     boxes.append(
                         BoundingBox(
                             x1=int(x1),
                             y1=int(y1),
                             x2=int(x2),
                             y2=int(y2),
-                            cls_id=int(self.CLS_MAP[int(cls_id)]),
+                            cls_id=int(mapped_cls_id),
                             conf=float(conf),
                         )
                     )
-                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)
-                # boxes = self.suppress_quasi_total_containment(boxes)
-                # boxes = self.suppress_small_contained(boxes)
-                # img_bgr = batch_images[frame_idx_in_batch]
-                # player_indices: List[int] = []
-                # player_feats: List[np.ndarray] = []
-                # for i, bb in enumerate(boxes):
-                #     if int(bb.cls_id) == 2:
-                #         hs = self._hs_feature_from_roi(img_bgr, bb)
-                #         player_indices.append(i)
-                #         player_feats.append(hs)
-                # cluster_centers = None
-                # n_players = len(player_feats)
-                # if n_players >= 2:
-                #     feats = np.vstack(player_feats)
-                #     labels, centers = self._assign_players_two_clusters(feats)
-                #     order = np.argsort(centers[:, 0])
-                #     centers = centers[order]
-                #     remap = {old_idx: new_idx for new_idx, old_idx in enumerate(order)}
-                #     labels = np.vectorize(remap.get)(labels)
-                #     cluster_centers = centers
-                #     for idx_in_list, lbl in zip(player_indices, labels):
-                #         boxes[idx_in_list].cls_id = 6 if int(lbl) == 0 else 7
-                # elif n_players == 1:
-                #     hue, _ = player_feats[0]
-                #     boxes[player_indices[0]].cls_id = 6 if float(hue) < self.SINGLE_PLAYER_HUE_PIVOT else 7
-                # self._reclass_extra_goalkeepers(img_bgr, boxes, cluster_centers)
-                bboxes[offset + frame_idx_in_batch] = boxes
+            # Handle footballs - keep only the best one
+            footballs = [bb for bb in boxes if int(bb.cls_id) == 0]
+            if len(footballs) > 1:
+                best_ball = max(footballs, key=lambda b: b.conf)
+                boxes = [bb for bb in boxes if int(bb.cls_id) != 0]
+                boxes.append(best_ball)
+        
+            bboxes[offset + frame_id] = boxes
+        return bboxes
+
+
+    def predict_batch(self, batch_images: List[ndarray], offset: int, n_keypoints: int) -> List[TVFrameResult]:        
+        start = time.time()
+        detection_results = self._detect_objects_batch(batch_images)
+        end = time.time()
+        print(f"Detection time: {end - start}")
+        start = time.time()
+        bboxes = self._team_classify(detection_results, batch_images, offset)
+        end = time.time()
+        print(f"Team classify time: {end - start}")
 
         pitch_batch_size = min(self.pitch_batch_size, len(batch_images))
         keypoints: Dict[int, List[Tuple[int, int]]] = {}
+
+        start = time.time()
         while True:
-            # try:
             gc.collect()
             if torch.cuda.is_available():
-                tf.keras.backend.clear_session()
                 torch.cuda.empty_cache()
                 torch.cuda.synchronize()
-            device_str = "cuda" if torch.cuda.is_available() else "cpu"
+            device_str = "cuda"
             keypoints_result = process_batch_input(
                 batch_images,
                 self.keypoints_model,
@@ -344,21 +405,10 @@ class Miner:
                     else:
                         frame_keypoints = frame_keypoints[:n_keypoints]
                     keypoints[offset + frame_number_in_batch] = frame_keypoints
-            print("✅ Keypoints predicted")
             break
-            # except RuntimeError as e:
-            #     print(self.pitch_batch_size)
-            #     print(e)
-            #     if "out of memory" in str(e):
-            #         if self.pitch_batch_size == 1:
-            #             break
-            #         self.pitch_batch_size = self.pitch_batch_size // 2 if self.pitch_batch_size > 1 else 1
-            #         pitch_batch_size = min(self.pitch_batch_size, len(batch_images))
-            #     else:
-            #         break
-            # except Exception as e:
-            #     print(f"❌ Error during keypoints prediction: {e}")
-            #     break
+        end = time.time()
+        print(f"Keypoint time: {end - start}")
+
 
         results: List[TVFrameResult] = []
         for frame_number in range(offset, offset + len(batch_images)):
@@ -373,7 +423,6 @@ class Miner:
 
         gc.collect()
         if torch.cuda.is_available():
-            tf.keras.backend.clear_session()
             torch.cuda.empty_cache()
             torch.cuda.synchronize()
 

pitch.py

@@ -660,28 +660,10 @@ def get_mapped_keypoints(kp_points):
             # mapped_points[key] = value
     return mapped_points
 
-def process_batch_input(frames, model, kp_threshold, device, batch_size=8):
+def process_batch_input(frames, model, kp_threshold, device, batch_size=16):
     """Process multiple input images in batch"""
     # Batch inference
     kp_results = inference_batch(frames, model, kp_threshold, device, batch_size)
     kp_results = [get_mapped_keypoints(kp) for kp in kp_results]
-    # Draw results and save
-    # for i, (frame, kp_points, input_path) in enumerate(zip(frames, kp_results, valid_paths)):
-    #     height, width = frame.shape[:2]
-        
-    #     # Apply mapping to get standard keypoint IDs
-    #     mapped_kp_points = get_mapped_keypoints(kp_points)
-        
-    #     for key, value in mapped_kp_points.items():
-    #         x = int(value['x'] * width)
-    #         y = int(value['y'] * height)
-    #         cv2.circle(frame, (x, y), 5, (0, 255, 0), -1)  # Green circles
-    #         cv2.putText(frame, str(key), (x+10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
-        
-    #     # Save result
-    #     output_path = input_path.replace('.png', '_result.png').replace('.jpg', '_result.jpg')
-    #     cv2.imwrite(output_path, frame)
-        
-    # print(f"Batch processing complete. Processed {len(frames)} images.")
 
     return kp_results