scorevision: push artifact

miner.py

@@ -27,22 +27,35 @@ class Miner:
     - Running batched predictions on images.
     - Parsing ML model outputs into structured results (TVFrameResult).
 
-    MODIFIED FOR TESTING: Uses standard yolov8n.pt and yolov8n-pose.pt
+    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.
+        Loads all ML models from the repository.
+        -----(Adjust as needed)----
+
+        Args:
+            path_hf_repo (Path):
+                Path to the downloaded HuggingFace Hub repository
+
+        Returns:
+            None
         """
-        # 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)")
+        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(
@@ -53,19 +66,31 @@ class Miner:
     ) -> 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]] = {}
-        # Run BBox prediction
-        bbox_model_results = self.bbox_model.predict(batch_images, verbose=False)
-        
+        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:
-                    # YOLOv8 standard output: x1, y1, x2, y2, conf, cls
                     x1, y1, x2, y2, conf, cls_id = box.tolist()
                     boxes.append(
                         BoundingBox(
@@ -81,33 +106,21 @@ class Miner:
         print("✅ BBoxes predicted")
 
         keypoints: dict[int, tuple[int, int]] = {}
-        # Run Pose/Keypoints prediction
-        keypoints_model_results = self.keypoints_model.predict(batch_images, verbose=False)
-        
+        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]] = []
-                
-                # 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
+                for part_points in detection.keypoints.data:
+                    for x, y, _ in part_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")