Upload miner.py

miner.py

@@ -13,11 +13,12 @@ class BoundingBox(BaseModel):
     cls_id: int
     conf: float
 
+
 class TVFrameResult(BaseModel):
     frame_id: int
     boxes: list[BoundingBox]
-    keypoints: list[tuple[int, int]] 
-    
+    keypoints: list[tuple[int, int]]
+
 
 class Miner:
     """
@@ -31,25 +32,24 @@ class 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:
+
+    def __init__(self, path_hf_repo: Path) -> None:
         """
         Loads all ML models from the repository.
         -----(Adjust as needed)----
 
         Args:
-            path_hf_repo (Path): 
+            path_hf_repo (Path):
                 Path to the downloaded HuggingFace Hub repository
 
         Returns:
-            None        
+            None
         """
-        self.bbox_model = YOLO(path_hf_repo/"football-player-detection.pt")
+        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")
+        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
@@ -65,21 +65,21 @@ class Miner:
         n_keypoints: int,
     ) -> list[TVFrameResult]:
         """
-        Miner prediction for a batch of images. 
+        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]): 
+            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. 
+            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): 
+            n_keypoints (int):
                 The number of keypoints expected for each frame in this challenge type.
 
         Returns:
-            list[TVFrameResult]: 
+            list[TVFrameResult]:
                 A list of predictions for each image in the batch
         """
 
@@ -112,6 +112,7 @@ class Miner:
                 if not hasattr(detection, "keypoints") or detection.keypoints is None:
                     continue
                 frame_keypoints: list[tuple[int, int]] = []
+                print(detection.keypoints.data)
                 for person_points in detection.keypoints.data:
                     for x, y in person_points:
                         frame_keypoints.append((int(x), int(y)))