scorevision: push artifact

miner.py

@@ -0,0 +1,325 @@
+from pathlib import Path
+
+from ultralytics import YOLO
+from numpy import ndarray
+from pydantic import BaseModel
+import numpy as np
+import cv2
+
+
+def calculate_iou(box1: dict, box2: dict) -> float:
+    x1 = max(box1["x1"], box2["x1"])
+    y1 = max(box1["y1"], box2["y1"])
+    x2 = min(box1["x2"], box2["x2"])
+    y2 = min(box1["y2"], box2["y2"])
+
+    intersection_width = max(0.0, x2 - x1)
+    intersection_height = max(0.0, y2 - y1)
+    intersection_area = intersection_width * intersection_height
+
+    box1_area = max(0.0, box1["x2"] - box1["x1"]) * max(
+        0.0, box1["y2"] - box1["y1"]
+    )
+    box2_area = max(0.0, box2["x2"] - box2["x1"]) * max(
+        0.0, box2["y2"] - box2["y1"]
+    )
+
+    union_area = box1_area + box2_area - intersection_area
+
+    if union_area <= 0:
+        return 0.0
+
+    return intersection_area / union_area
+
+
+def fuse_box_cluster(cluster: list[dict]) -> dict:
+    """
+    Merge box coordinates using confidence-weighted averaging.
+
+    The merged confidence is the highest confidence in the cluster.
+    """
+    weights = np.array(
+        [max(box["conf"], 1e-6) for box in cluster],
+        dtype=np.float32,
+    )
+
+    coordinates = np.array(
+        [
+            [box["x1"], box["y1"], box["x2"], box["y2"]]
+            for box in cluster
+        ],
+        dtype=np.float32,
+    )
+
+    fused_coordinates = np.average(
+        coordinates,
+        axis=0,
+        weights=weights,
+    )
+
+    best_box = max(cluster, key=lambda box: box["conf"])
+
+    return {
+        "x1": int(round(fused_coordinates[0])),
+        "y1": int(round(fused_coordinates[1])),
+        "x2": int(round(fused_coordinates[2])),
+        "y2": int(round(fused_coordinates[3])),
+        "cls_id": int(best_box["cls_id"]),
+        "conf": float(max(box["conf"] for box in cluster)),
+    }
+
+
+def merge_same_class_boxes(
+    detections: list[dict],
+    iou_threshold: float = 0.65,
+) -> list[dict]:
+    """
+    Merge overlapping boxes only when they have the same class ID.
+    """
+    merged_detections = []
+
+    class_ids = sorted({box["cls_id"] for box in detections})
+
+    for class_id in class_ids:
+        remaining = [
+            box.copy()
+            for box in detections
+            if box["cls_id"] == class_id
+        ]
+
+        remaining.sort(key=lambda box: box["conf"], reverse=True)
+
+        while remaining:
+            cluster = [remaining.pop(0)]
+
+            # Recalculate the fused box as new members join the cluster.
+            changed = True
+
+            while changed:
+                changed = False
+                current_fused_box = fuse_box_cluster(cluster)
+                not_merged = []
+
+                for candidate in remaining:
+                    iou = calculate_iou(current_fused_box, candidate)
+
+                    if iou >= iou_threshold:
+                        cluster.append(candidate)
+                        changed = True
+                    else:
+                        not_merged.append(candidate)
+
+                remaining = not_merged
+
+            merged_detections.append(fuse_box_cluster(cluster))
+
+    merged_detections.sort(
+        key=lambda box: box["conf"],
+        reverse=True,
+    )
+
+    return merged_detections
+
+
+def reference_model(image: np.ndarray, model: YOLO, confidence_threshold: float = 0.1, is_flipped: bool = False) -> list[dict]:
+    results = model.predict(
+        source=image,
+        conf=confidence_threshold,
+        verbose=False,
+    )
+
+    result = results[0]
+    detections = []
+
+    if result.boxes is not None:
+        coordinates = result.boxes.xyxy.cpu().tolist()
+        confidences = result.boxes.conf.cpu().tolist()
+        class_ids = result.boxes.cls.int().cpu().tolist()
+
+        for coordinates_xyxy, confidence, class_id in zip(
+            coordinates,
+            confidences,
+            class_ids,
+        ):
+            x1, y1, x2, y2 = coordinates_xyxy
+            
+            if is_flipped:
+                # Flip the x-coordinates for the horizontally flipped image
+                x1, x2 = image.shape[1] - x2, image.shape[1] - x1
+            detection = {
+                "x1": float(x1),
+                "y1": float(y1),
+                "x2": float(x2),
+                "y2": float(y2),
+                "cls_id": int(class_id),
+                "conf": float(confidence),
+            }
+            detections.append(detection)
+
+    return detections
+
+
+class BoundingBox(BaseModel):
+    x1: int
+    y1: int
+    x2: int
+    y2: int
+    cls_id: int
+    conf: float
+
+
+class Polygon(BaseModel):
+    cls_id: int
+    conf: float
+    points: list[tuple[int, int]]
+
+
+class TVFrameResult(BaseModel):
+    frame_id: int
+    boxes: list[BoundingBox] | None = None
+    polygons: list[Polygon] | None = None
+    keypoints: list[tuple[int, int]] | None = None
+
+
+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 / "weights.onnx")
+        print(f"✅ BBox Model Loaded")
+
+        self.model_class_names = ["balaclava", "bat", "glove", "graffiti", "hoodie", "spray paint"]
+        self.subnet_class_names = ["balaclava", "hoodie", "glove", "bat", "spray paint", "graffiti"]
+
+        self._class_id_map = np.fromiter(
+            map(self.subnet_class_names.index, self.model_class_names), dtype=np.int32,
+            count=len(self.model_class_names),
+        )
+
+        self._conf_thres_array = np.array(
+            [0.44, 0.65, 0.55, 0.25, 0.45, 0.26], dtype=np.float32,
+        )
+        self._bonus_array = np.array(
+            [0.23, 0.3, 0.45, 0.14, 0.45, 0.41], dtype=np.float32,
+        )
+
+    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"Model: {type(self.bbox_model).__name__}"
+
+    def predict_crime(self, image: np.ndarray, model: YOLO, confidence_threshold: float = 0.1) -> list[dict]:
+        flipped = cv2.flip(image, 1)  # Flip the image horizontally
+
+        original_detections = reference_model(
+            image=image,
+            model=model,
+            confidence_threshold=confidence_threshold,
+            is_flipped=False,
+        )
+        flipped_detections = reference_model(
+            image=flipped,
+            model=model,
+            confidence_threshold=confidence_threshold,
+            is_flipped=True,
+        )
+
+        detections = original_detections + flipped_detections
+
+        merged_boxes = merge_same_class_boxes(
+            detections,
+            iou_threshold=0.65,
+        )
+
+        return merged_boxes
+
+    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]] = {}
+
+        for frame_number_in_batch, image in enumerate(batch_images):
+            merged_boxes = self.predict_crime(
+                image=image,
+                model=self.bbox_model,
+                confidence_threshold=0.1,
+            )
+
+            boxes = []
+            for box in merged_boxes:
+                x1, y1, x2, y2, model_cls_id, model_conf = box.values()
+                sn_cls_id = self._class_id_map[int(model_cls_id)]
+
+                if model_conf < self._conf_thres_array[sn_cls_id]:
+                    continue
+                
+                bossted_conf = min(model_conf + self._bonus_array[sn_cls_id], 0.99)
+                boxes.append(
+                    BoundingBox(
+                        x1=int(x1),
+                        y1=int(y1),
+                        x2=int(x2),
+                        y2=int(y2),
+                        cls_id=int(sn_cls_id),
+                        conf=float(bossted_conf),
+                    )
+                )
+            bboxes[offset + frame_number_in_batch] = boxes
+        print("✅ BBoxes 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, []),
+                )
+            )
+        print("✅ Combined results as TVFrameResult")
+        return results