Upload miner.py

miner.py

@@ -0,0 +1,139 @@
+from pathlib import Path
+
+from ultralytics import YOLO
+from numpy import ndarray
+from pydantic import BaseModel
+
+
+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:
+    """
+    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
+    """
+    
+    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        
+        """
+        self.bbox_model = YOLO(path_hf_repo/"football-player-detection.pt")
+        print(f"✅ BBox Model Loaded")
+        self.keypoints_model = YOLO(path_hf_repo/"football-pitch-detection.pt")
+        print(f"✅ Keypoints Model Loaded")
+
+
+    def __repr__(self) -> str:
+        """
+        Information about miner returned in the health endpoint
+        to inspect the loaded ML models (and their types)
+        -----(Adjust as needed)----
+        """
+        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. 
+        Handles the orchestration of ML models and any preprocessing and postprocessing
+        -----(Adjust as needed)----
+
+        Args:
+            batch_images (list[np.ndarray]): 
+                A list of images (as NumPy arrays) to process in this batch.
+            offset (int): 
+                The frame number corresponding to the first image in the batch. 
+                Used to correctly index frames in the output results.
+            n_keypoints (int): 
+                The number of keypoints expected for each frame in this challenge type.
+
+        Returns:
+            list[TVFrameResult]: 
+                A list of predictions for each image in the batch
+        """
+
+        bboxes: dict[int, list[BoundingBox]] = {}
+        bbox_model_results = self.bbox_model.predict(batch_images)
+        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:
+                    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]] = {}
+        keypoints_model_results = self.keypoints_model.predict(batch_images)
+        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]] = []
+                for person_points in detection.keypoints.data:
+                    for x, y in person_points:
+                        frame_keypoints.append((int(x), int(y)))
+                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