scorevision: push artifact

miner.py

@@ -1,8 +1,18 @@
 from pathlib import Path
+from typing import List, Tuple, Dict
+import sys
+import os
 
-from ultralytics import YOLO
 from numpy import ndarray
 from pydantic import BaseModel
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+from team_cluster import TeamClassifier
+from utils import (
+    BoundingBox, 
+    Constants,
+)
+from inference import predict_batch, load_model
 
 
 class BoundingBox(BaseModel):
@@ -16,111 +26,32 @@ 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).
-
-    MODIFIED FOR TESTING: Uses standard yolov8n.pt and yolov8n-pose.pt
-    """
+    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:
-        """
-        Loads all ML models.
-        """
-        # Using standard YOLOv8 nano models that will be automatically downloaded
-        # if not present. This avoids the need for custom .pt files for testing.
-        self.bbox_model = YOLO("yolov8n.pt")
-        print(f"✅ BBox Model Loaded (yolov8n)")
-        
-        self.keypoints_model = YOLO("yolov8n-pose.pt")
-        print(f"✅ Keypoints Model Loaded (yolov8n-pose)")
+        print("model laoding")
+        self.health = load_model(path_hf_repo)
 
     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]:
-        """
-        Miner prediction for a batch of images.
-        """
-
-        bboxes: dict[int, list[BoundingBox]] = {}
-        # Run BBox prediction
-        bbox_model_results = self.bbox_model.predict(batch_images, verbose=False)
-        
-        if bbox_model_results is not None:
-            for frame_number_in_batch, detection in enumerate(bbox_model_results):
-                if not hasattr(detection, "boxes") or detection.boxes is None:
-                    continue
-                boxes = []
-                for box in detection.boxes.data:
-                    # YOLOv8 standard output: x1, y1, x2, y2, conf, cls
-                    x1, y1, x2, y2, conf, cls_id = box.tolist()
-                    boxes.append(
-                        BoundingBox(
-                            x1=int(x1),
-                            y1=int(y1),
-                            x2=int(x2),
-                            y2=int(y2),
-                            cls_id=int(cls_id),
-                            conf=float(conf),
-                        )
-                    )
-                bboxes[offset + frame_number_in_batch] = boxes
-        print("✅ BBoxes predicted")
-
-        keypoints: dict[int, tuple[int, int]] = {}
-        # Run Pose/Keypoints prediction
-        keypoints_model_results = self.keypoints_model.predict(batch_images, verbose=False)
-        
-        if keypoints_model_results is not None:
-            for frame_number_in_batch, detection in enumerate(keypoints_model_results):
-                if not hasattr(detection, "keypoints") or detection.keypoints is None:
-                    continue
-                
-                frame_keypoints: list[tuple[int, int]] = []
-                
-                # Check if keypoints data exists and has the expected shape/content
-                if detection.keypoints.data is not None and len(detection.keypoints.data) > 0:
-                     # Taking the first person detected for keypoints (simplification for testing)
-                     # YOLO pose output is typically [num_people, num_kpts, 3] (x, y, conf)
-                    first_person_kpts = detection.keypoints.data[0]
-                    for kpt in first_person_kpts:
-                       x, y = kpt[0], kpt[1] # extracting x, y
-                       frame_keypoints.append((int(x), int(y)))
-                
-                # Padding or truncating to match expected n_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")
-
-        results: list[TVFrameResult] = []
-        for frame_number in range(offset, offset + len(batch_images)):
-            results.append(
-                TVFrameResult(
-                    frame_id=frame_number,
-                    boxes=bboxes.get(frame_number, []),
-                    keypoints=keypoints.get(
-                        frame_number, [(0, 0) for _ in range(n_keypoints)]
-                    ),
-                )
-            )
-        print("✅ Combined results as TVFrameResult")
-        return results
+        return self.health
+
+    def predict_batch(self, batch_images: List[ndarray], offset: int, n_keypoints: int) -> List[TVFrameResult]:    
+        results = predict_batch(
+            batch_images,
+            offset,
+            n_keypoints
+        )
+        return results