Upload folder using huggingface_hub

README.md

@@ -1,92 +1,92 @@
-# 🚀 Example Chute for Turbovision 🪂
-
-This repository demonstrates how to deploy a **Chute** via the **Turbovision CLI**, hosted on **Hugging Face Hub**.
-It serves as a minimal example showcasing the required structure and workflow for integrating machine learning models, preprocessing, and orchestration into a reproducible Chute environment.
-
-## Repository Structure 
-The following two files **must be present** (in their current locations) for a successful deployment — their content can be modified as needed:
-
-| File | Purpose |
-|------|----------|
-| `miner.py` | Defines the ML model type(s), orchestration, and all pre/postprocessing logic. |
-| `config.yml` | Specifies machine configuration (e.g., GPU type, memory, environment variables). |
-
-Other files — e.g., model weights, utility scripts, or dependencies — are **optional** and can be included as needed for your model. Note: Any required assets must be defined or contained **within this repo**, which is fully open-source, since all network-related operations (downloading challenge data, weights, etc.) are disabled **inside the Chute** 
-
-## Overview
-
-Below is a high-level diagram showing the interaction between Huggingface, Chutes and Turbovision:
-
-![](../images/miner.png)
-
-## Local Testing
-After editing the `config.yml` and `miner.py` and saving it into your Huggingface Repo, you will want to test it works locally. 
-
-1. Copy the file `scorevision/chute_tmeplate/turbovision_chute.py.j2` as a python file called `my_chute.py` and fill in the missing variables:
-```python
-HF_REPO_NAME = "{{ huggingface_repository_name }}"
-HF_REPO_REVISION = "{{ huggingface_repository_revision }}"
-CHUTES_USERNAME = "{{ chute_username }}"
-CHUTE_NAME = "{{ chute_name }}"
-```
-
-2. Run the following command to build the chute locally (Caution: there are known issues with the docker location when running this on a mac) 
-```bash
-chutes build my_chute:chute --local --public
-```
-
-3. Run the name of the docker image just built (i.e. `CHUTE_NAME`) and enter it
-```bash
-docker run -p 8000:8000 -e CHUTES_EXECUTION_CONTEXT=REMOTE -it <image-name> /bin/bash
-```
-
-4. Run the file from within the container
-```bash
-chutes run my_chute:chute --dev --debug
-```
-
-5. In another terminal, test the local endpoints to ensure there are no bugs
-```bash
-curl -X POST http://localhost:8000/health -d '{}'
-curl -X POST http://localhost:8000/predict -d '{"url": "https://scoredata.me/2025_03_14/35ae7a/h1_0f2ca0.mp4","meta": {}}'
-```
-
-## Live Testing
-1. If you have any chute with the same name (ie from a previous deployment), ensure you delete that first (or you will get an error when trying to build).
-```bash
-chutes chutes list
-```
-Take note of the chute id that you wish to delete (if any)
-```bash
-chutes chutes delete <chute-id>
-```
-
-You should also delete its associated image 
-```bash
-chutes images list
-```
-Take note of the chute image id
-```bash
-chutes images delete <chute-image-id>
-```
-
-2. Use Turbovision's CLI to build, deploy and commit on-chain (Note: you can skip the on-chain commit using `--no-commit`.  You can also specify a past huggingface revision to point to using `--revision` and/or the local files you want to upload to your huggingface repo using `--model-path`)
-```bash
-sv -vv push
-```
-
-3. When completed, warm up the chute (if its cold 🧊). (You can confirm its status using `chutes chutes list` or `chutes chutes get <chute-id>` if you already know its id). Note: Warming up can sometimes take a while but if the chute runs without errors (should be if you've tested locally first) and there are sufficient nodes (i.e. machines) available matching the `config.yml` you specified, the chute should become hot 🔥!
-```bash
-chutes warmup <chute-id>
-```
-
-4. Test the chute's endpoints
-```bash
-curl -X POST https://<YOUR-CHUTE-SLUG>.chutes.ai/health -d '{}' -H "Authorization: Bearer $CHUTES_API_KEY"
-curl -X POST https://<YOUR-CHUTE-SLUG>.chutes.ai/predict -d '{"url": "https://scoredata.me/2025_03_14/35ae7a/h1_0f2ca0.mp4","meta": {}}' -H "Authorization: Bearer $CHUTES_API_KEY"
-```
-
-5. Test what your chute would get on a validator (this also applies any validation/integrity checks which may fail if you did not use the Turbovision CLI above to deploy the chute)
-```bash
-sv -vv run-once
+# 🚀 Example Chute for Turbovision 🪂
+
+This repository demonstrates how to deploy a **Chute** via the **Turbovision CLI**, hosted on **Hugging Face Hub**.
+It serves as a minimal example showcasing the required structure and workflow for integrating machine learning models, preprocessing, and orchestration into a reproducible Chute environment.
+
+## Repository Structure 
+The following two files **must be present** (in their current locations) for a successful deployment — their content can be modified as needed:
+
+| File | Purpose |
+|------|----------|
+| `miner.py` | Defines the ML model type(s), orchestration, and all pre/postprocessing logic. |
+| `config.yml` | Specifies machine configuration (e.g., GPU type, memory, environment variables). |
+
+Other files — e.g., model weights, utility scripts, or dependencies — are **optional** and can be included as needed for your model. Note: Any required assets must be defined or contained **within this repo**, which is fully open-source, since all network-related operations (downloading challenge data, weights, etc.) are disabled **inside the Chute** 
+
+## Overview
+
+Below is a high-level diagram showing the interaction between Huggingface, Chutes and Turbovision:
+
+![](../images/miner.png)
+
+## Local Testing
+After editing the `config.yml` and `miner.py` and saving it into your Huggingface Repo, you will want to test it works locally. 
+
+1. Copy the file `scorevision/chute_tmeplate/turbovision_chute.py.j2` as a python file called `my_chute.py` and fill in the missing variables:
+```python
+HF_REPO_NAME = "{{ huggingface_repository_name }}"
+HF_REPO_REVISION = "{{ huggingface_repository_revision }}"
+CHUTES_USERNAME = "{{ chute_username }}"
+CHUTE_NAME = "{{ chute_name }}"
+```
+
+2. Run the following command to build the chute locally (Caution: there are known issues with the docker location when running this on a mac) 
+```bash
+chutes build my_chute:chute --local --public
+```
+
+3. Run the name of the docker image just built (i.e. `CHUTE_NAME`) and enter it
+```bash
+docker run -p 8000:8000 -e CHUTES_EXECUTION_CONTEXT=REMOTE -it <image-name> /bin/bash
+```
+
+4. Run the file from within the container
+```bash
+chutes run my_chute:chute --dev --debug
+```
+
+5. In another terminal, test the local endpoints to ensure there are no bugs
+```bash
+curl -X POST http://localhost:8000/health -d '{}'
+curl -X POST http://localhost:8000/predict -d '{"url": "https://scoredata.me/2025_03_14/35ae7a/h1_0f2ca0.mp4","meta": {}}'
+```
+
+## Live Testing
+1. If you have any chute with the same name (ie from a previous deployment), ensure you delete that first (or you will get an error when trying to build).
+```bash
+chutes chutes list
+```
+Take note of the chute id that you wish to delete (if any)
+```bash
+chutes chutes delete <chute-id>
+```
+
+You should also delete its associated image 
+```bash
+chutes images list
+```
+Take note of the chute image id
+```bash
+chutes images delete <chute-image-id>
+```
+
+2. Use Turbovision's CLI to build, deploy and commit on-chain (Note: you can skip the on-chain commit using `--no-commit`.  You can also specify a past huggingface revision to point to using `--revision` and/or the local files you want to upload to your huggingface repo using `--model-path`)
+```bash
+sv -vv push
+```
+
+3. When completed, warm up the chute (if its cold 🧊). (You can confirm its status using `chutes chutes list` or `chutes chutes get <chute-id>` if you already know its id). Note: Warming up can sometimes take a while but if the chute runs without errors (should be if you've tested locally first) and there are sufficient nodes (i.e. machines) available matching the `config.yml` you specified, the chute should become hot 🔥!
+```bash
+chutes warmup <chute-id>
+```
+
+4. Test the chute's endpoints
+```bash
+curl -X POST https://<YOUR-CHUTE-SLUG>.chutes.ai/health -d '{}' -H "Authorization: Bearer $CHUTES_API_KEY"
+curl -X POST https://<YOUR-CHUTE-SLUG>.chutes.ai/predict -d '{"url": "https://scoredata.me/2025_03_14/35ae7a/h1_0f2ca0.mp4","meta": {}}' -H "Authorization: Bearer $CHUTES_API_KEY"
+```
+
+5. Test what your chute would get on a validator (this also applies any validation/integrity checks which may fail if you did not use the Turbovision CLI above to deploy the chute)
+```bash
+sv -vv run-once
 ```

config.yml

@@ -2,40 +2,28 @@ Image:
   from_base: parachutes/python:3.12
   run_command:
     - pip install --upgrade setuptools wheel
-    - pip install huggingface_hub==0.19.4 requests opencv-python-headless pydantic onnxruntime onnxruntime-gpu scikit-learn tensorflow torch-tensorrt==2.7 torch==2.7.1 torchvision==0.22.1 pyyaml
-
+    - 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"
   set_workdir: /app
 
+
 NodeSelector:
   gpu_count: 1
-  min_vram_gb_per_gpu: 24
-  # include:
-  #   - a100
-  #   - a100_40gb
-  #   - "3090"
-  #   - a40
-  #   - a6000
-  exclude:
-    - h100
+  min_vram_gb_per_gpu: 16
+  include:
     - a100
+    - "3090"
+    - a40
+    - a6000
     - l40s
-    - mi300x
+  exclude:
     - b200
     - h200
     - h20
-    - h800
-    - h100_sxm
-    - h100_nvl
-    - a100_sxm
-    - a100_40gb_sxm
-    - a100_40gb
-    - l40
-    - pro_6000
-    - a6000_ada
-    - '5090'
+    - mi300x
 
 Chute:
-  timeout_seconds: 300
-  concurrency: 4  # Reduced concurrency to limit memory usage
-  max_instances: 5  # Reduced max instances to limit memory usage
-  scaling_threshold: 0.5  # Higher threshold to reduce scaling frequency
+  timeout_seconds: 900
+  concurrency: 4
+  max_instances: 5
+  scaling_threshold: 0.5

miner.py

@@ -1,47 +1,43 @@
 from pathlib import Path
+from typing import List, Tuple, Dict
+import sys
+import os
 
 from numpy import ndarray
 import numpy as np
 from pydantic import BaseModel
-import sys, os
+import cv2
+
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+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'
-# Suppress ONNX Runtime warnings
-os.environ['ORT_LOGGING_LEVEL'] = '3'
+os.environ["CUDA_LAUNCH_BLOCKING"] = "0"
+os.environ["ORT_LOGGING_LEVEL"] = "3"
+os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
 import logging
-logging.getLogger('tensorflow').setLevel(logging.ERROR)
-
 import tensorflow as tf
+from tensorflow.keras import mixed_precision
+import torch._dynamo
+import torch
+import torch_tensorrt
+import gc
+from ultralytics import YOLO
+from pitch import process_batch_input
+
+logging.getLogger("tensorflow").setLevel(logging.ERROR)
 tf.config.threading.set_intra_op_parallelism_threads(16)
 tf.config.threading.set_inter_op_parallelism_threads(2)
-os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
 tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
-tf.get_logger().setLevel('ERROR')
+tf.get_logger().setLevel("ERROR")
 tf.autograph.set_verbosity(0)
-
-from tensorflow.keras import mixed_precision
-mixed_precision.set_global_policy('mixed_float16')
+mixed_precision.set_global_policy("mixed_float16")
 tf.config.optimizer.set_jit(True)
-
-import torch._dynamo
 torch._dynamo.config.suppress_errors = True
-import onnxruntime as ort
-import gc
-
-import torch
-import torch_tensorrt
-import torchvision.transforms as T
-import yaml
-import cv2
 
-from player import player_detection_result
-from pitch import process_batch_input, get_cls_net, get_cls_net_l
 
 class BoundingBox(BaseModel):
     x1: int
@@ -54,57 +50,21 @@ class BoundingBox(BaseModel):
 
 class TVFrameResult(BaseModel):
     frame_id: int
-    boxes: list[BoundingBox]
-    keypoints: list[tuple[int, int]]
+    boxes: List[BoundingBox]
+    keypoints: List[Tuple[int, int]]
 
-class Miner:
-    """
-    This class is responsible for:
-    - Loading ML models.
-    - Running batched predictions on images.
-    - Parsing ML model outputs into structured results (TVFrameResult).
 
-    This class can be modified, but it must have the following to be compatible with the chute:
-        - be named `Miner`
-        - have a `predict_batch` function with the inputs and outputs specified
-        - be stored in a file called `miner.py` which lives in the root of the HFHub repo
-    """
+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
 
     def __init__(self, path_hf_repo: Path) -> None:
-        """
-        Loads all ML models from the repository.
-        -----(Adjust as needed)----
-
-        Args:
-            path_hf_repo (Path):
-                Path to the downloaded HuggingFace Hub repository
-
-        Returns:
-            None
-        """
-        global torch
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-        providers = [
-            'CUDAExecutionProvider',
-            'CPUExecutionProvider'
-            
-        ]
-        # providers = [ 'CPUExecutionProvider']
-        model_path = path_hf_repo / "object-detection.onnx"
-        session = ort.InferenceSession(model_path, providers=providers)
-        input_name = session.get_inputs()[0].name
-        height = width = 640
-        dummy = np.zeros((1, 3, height, width), dtype=np.float32)
-        session.run(None, {input_name: dummy})
-        model = session
-        self.bbox_model = model
-        print(f"✅ BBox Model Loaded")
-
-        self.kp_threshold = 0.1
-        # self.lp_threshold = 0.7
-
-        model_kp_path = path_hf_repo / 'SV_kp.engine'
+        self.bbox_model = YOLO(path_hf_repo / "player.pt")
+        print(" BBox Model (objdetect.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)
 
         @torch.inference_mode()
@@ -114,114 +74,214 @@ class Miner:
             return output
 
         run_inference(model_kp, torch.randn(8, 3, 540, 960, device=device, dtype=torch.float32))
-        # model_kp_path = path_hf_repo / 'SV_kp'
-        # model_lp_path = path_hf_repo / 'SV_lines'
-        # config_kp_path = path_hf_repo / 'hrnetv2_w48.yaml'
-        # config_lp_path = path_hf_repo / 'hrnetv2_w48_l.yaml'
-        # cfg_kp = yaml.safe_load(open(config_kp_path, 'r'))
-        # cfg_lp = yaml.safe_load(open(config_lp_path, 'r'))
-
-        # loaded_state_kp = torch.load(model_kp_path, map_location=device)
-        # model_kp = get_cls_net(cfg_kp)
-        # model_kp.load_state_dict(loaded_state_kp)
-        # model_kp.to(device)
-        # model_kp.eval()
-          
-        # loaded_state_lp = torch.load(model_lp_path, map_location=device)
-        # model_lp = get_cls_net_l(cfg_lp)
-        # model_lp.load_state_dict(loaded_state_lp)
-        # model_lp.to(device)
-        # model_lp.eval()
-
-        # self.transform = T.Resize((540, 960))
-
         self.keypoints_model = model_kp
-        # self.lines_model = model_lp
-        
-        # print("🔥 Warming up compiled models...")
-        # self._warmup_models(device)
-        
-        # Increase batch sizes for better GPU utilization
-        self.player_batch_size = 16  # Increased from 32
-        self.pitch_batch_size = 8   # Increased from 32
-        print(f"✅ Keypoints Model Loaded")
+        self.kp_threshold = 0.1
+        self.pitch_batch_size = 8
+        print("✅ Keypoints Model Loaded")
 
     def __repr__(self) -> str:
-        return f"BBox Model: {type(self.bbox_model).__name__}\nKeypoints Model: {type(self.keypoints_model).__name__}"
-
-    def predict_batch(
-        self,
-        batch_images: list[ndarray],
-        offset: int,
-        n_keypoints: int,
-    ) -> list[TVFrameResult]:
-        player_batch_size = min(self.player_batch_size, len(batch_images))
-        bboxes: dict[int, list[BoundingBox]] = {}
-        while True:
-            try:
-                gc.collect()
-                if torch.cuda.is_available():
-                    tf.keras.backend.clear_session()
-                    torch.cuda.empty_cache()
-                    torch.cuda.synchronize()
-
-                bbox_model_results, _, _, _ = player_detection_result(batch_images, player_batch_size, self.bbox_model)
-                if bbox_model_results is not None and len(bbox_model_results) > 0:
-                    for frame_number_in_batch, detections in enumerate(bbox_model_results):
-                        # Ensure frame_number_in_batch is within batch_images bounds
-                        if frame_number_in_batch >= len(batch_images):
-                            print(f"⚠️  Warning: bbox_model_results has more frames ({len(bbox_model_results)}) than batch_images ({len(batch_images)}). Skipping extra frames.")
+        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
+        else:
+            h_mean, s_mean = self._mean_hs(roi)
+        return np.array([h_mean, s_mean], dtype=np.float32)
+
+    def _ioa(self, a: BoundingBox, b: BoundingBox) -> float:
+        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]:
+        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]:
+                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
-                        
-                        boxes = []
-                        if detections is not None and isinstance(detections, (list, tuple)):
-                            for detection in detections:
-                                try:
-                                    # Detection format from player.py: {"id": int, "bbox": [x1, y1, x2, y2], "class_id": int}
-                                    if isinstance(detection, dict):
-                                        x1, y1, x2, y2 = detection.get("bbox", [0, 0, 0, 0])
-                                        cls_id = detection.get("class_id", 0)
-                                        conf = detection.get("conf", 0.0)
-                                    else:
-                                        # Handle tuple/array format: (box, score, cls)
-                                        if len(detection) >= 3:
-                                            x1, y1, x2, y2 = detection[0] if hasattr(detection[0], '__iter__') else [0, 0, 0, 0]
-                                            conf = detection[1] if len(detection) > 1 else 0.0
-                                            cls_id = detection[2] if len(detection) > 2 else 0
-                                        else:
-                                            continue
-                                    
-                                    boxes.append(
-                                        BoundingBox(
-                                            x1=int(x1),
-                                            y1=int(y1),
-                                            x2=int(x2),
-                                            y2=int(y2),
-                                            cls_id=int(cls_id),
-                                            conf=float(conf),
-                                        )
-                                    )
-                                except (KeyError, TypeError, ValueError, IndexError) as det_err:
-                                    print(f"⚠️  Warning: Could not parse detection: {det_err}")
-                                    continue
-                        bboxes[offset + frame_number_in_batch] = boxes
-                print("✅ BBoxes predicted")
-                break
-            except RuntimeError as e:
-                print(self.player_batch_size)
-                if 'out of memory' in str(e):
-                    if self.player_batch_size == 1:
-                        break
-                    self.player_batch_size = self.player_batch_size // 2 if self.player_batch_size > 1 else 1
-                    player_batch_size = min(self.player_batch_size, len(batch_images))
                 else:
-                    break
-            except Exception as e:
-                print(f"❌ Error during bbox prediction: {e}")
-                break
+                    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:
+                    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
+                    boxes.append(
+                        BoundingBox(
+                            x1=int(x1),
+                            y1=int(y1),
+                            x2=int(x2),
+                            y2=int(y2),
+                            cls_id=int(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
 
         pitch_batch_size = min(self.pitch_batch_size, len(batch_images))
-        keypoints: dict[int, list[tuple[int, int]]] = {}
+        keypoints: Dict[int, List[Tuple[int, int]]] = {}
         while True:
             try:
                 gc.collect()
@@ -229,28 +289,21 @@ class Miner:
                     tf.keras.backend.clear_session()
                     torch.cuda.empty_cache()
                     torch.cuda.synchronize()
-                # Removed expensive memory clearing operations for speed
+                device_str = "cuda" if torch.cuda.is_available() else "cpu"
                 keypoints_result = process_batch_input(
                     batch_images,
                     self.keypoints_model,
                     self.kp_threshold,
-                    'cuda' if torch.cuda.is_available() else 'cpu',
-                    batch_size=pitch_batch_size
+                    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):
-                        # Ensure frame_number_in_batch is within batch_images bounds
                         if frame_number_in_batch >= len(batch_images):
-                            print(f"⚠️  Warning: keypoints_result has more frames ({len(keypoints_result)}) than batch_images ({len(batch_images)}). Skipping extra frames.")
                             break
-                        
-                        frame_keypoints: list[tuple[int, int]] = []
-                        
-                        # Get image dimensions for conversion from normalized to pixel coordinates
+                        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
@@ -258,31 +311,24 @@ class Miner:
                                     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:
-                                                # Convert normalized coordinates to pixel coordinates
-                                                x = int(kp_data['x'] * width)
-                                                y = int(kp_data['y'] * height)
-                                        except (KeyError, TypeError, ValueError) as kp_err:
-                                            print(f"⚠️  Warning: Could not parse keypoint {kp_idx}: {kp_err}")
+                                            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) as img_err:
-                            print(f"⚠️  Warning: Could not process frame {frame_number_in_batch}: {img_err}")
-                            # Create default keypoints if processing fails
+                        except (IndexError, ValueError, AttributeError):
                             frame_keypoints = [(0, 0)] * 32
-                        
-                        # Pad or truncate to match expected number of keypoints
                         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
-
                 print("✅ Keypoints predicted")
                 break
             except RuntimeError as e:
                 print(self.pitch_batch_size)
-                if 'out of memory' in str(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
@@ -293,13 +339,10 @@ class Miner:
                 print(f"❌ Error during keypoints prediction: {e}")
                 break
 
-        # Combine results
-        results: list[TVFrameResult] = []
-        for i, frame_number in enumerate(range(offset, offset + len(batch_images))):
+        results: List[TVFrameResult] = []
+        for frame_number in range(offset, offset + len(batch_images)):
             frame_boxes = bboxes.get(frame_number, [])
             frame_keypoints = keypoints.get(frame_number, [(0, 0) for _ in range(n_keypoints)])
-            
-            # Create result object
             result = TVFrameResult(
                 frame_id=frame_number,
                 boxes=frame_boxes,
@@ -307,12 +350,10 @@ class Miner:
             )
             results.append(result)
 
-        print("✅ Combined results as TVFrameResult")
-        
         gc.collect()
         if torch.cuda.is_available():
             tf.keras.backend.clear_session()
             torch.cuda.empty_cache()
             torch.cuda.synchronize()
-            
-        return results
+
+        return results

objdetect.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8bbacfcb38e38b1b8816788e9e6e845160533719a0b87b693d58b932380d0d28
+size 152961687

player.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce9fc31f61e6f156f786077abb8eef36b0836bda1ef07d1d0ba82d43ae0ecd0b
+size 22540152