Add files using upload-large-folder tool

config.yml

@@ -0,0 +1,21 @@
+Image:
+  from_base: parachutes/python:3.12
+  run_command:
+    - pip install --upgrade setuptools wheel
+    - pip install ultralytics==8.3.222 torch==2.9.0 opencv-python-headless numpy pydantic
+  set_workdir: /app
+
+NodeSelector:
+  gpu_count: 1
+  min_vram_gb_per_gpu: 16
+  exclude:
+    - "5090"
+    - b200
+    - h200
+    - mi300x
+
+Chute:
+  timeout_seconds: 300
+  concurrency: 4
+  max_instances: 5
+  scaling_threshold: 0.5

keypointdetect.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6dd10dba85895c92760cdb5a99c5cfca899c68f361a66c5448f38a187280ee1f
+size 6849672

miner.py

@@ -0,0 +1,296 @@
+from pathlib import Path
+from typing import List, Tuple, Dict, Optional
+
+from ultralytics import YOLO
+from numpy import ndarray
+from pydantic import BaseModel
+import numpy as np
+import cv2
+
+
+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:
+    QUASI_TOTAL_IOA: float = 0.90
+    SMALL_CONTAINED_IOA: float = 0.85
+    SMALL_RATIO_MAX: float = 0.50
+    SINGLE_PLAYER_HUE_PIVOT: float = 90.0
+    CORNER_INDICES = {0, 5, 24, 29}
+
+    def __init__(self, path_hf_repo: Path) -> None:
+        self.bbox_model = YOLO(path_hf_repo / "objdetect.pt")
+        print("BBox Model (objdetect.pt) Loaded")
+        self.keypoints_model = YOLO(path_hf_repo / "keypointdetect.pt")
+        print("Keypoints Model (keypointdetect.pt) Loaded")
+
+    def __repr__(self) -> str:
+        return (
+            f"BBox Model: {type(self.bbox_model).__name__}\n"
+            f"Keypoints Model: {type(self.keypoints_model).__name__}"
+        )
+
+    @staticmethod
+    def _clip_box_to_image(x1: int, y1: int, x2: int, y2: int, w: int, h: int) -> Tuple[int, int, int, int]:
+        x1 = max(0, min(int(x1), w - 1))
+        y1 = max(0, min(int(y1), h - 1))
+        x2 = max(0, min(int(x2), w - 1))
+        y2 = max(0, min(int(y2), h - 1))
+        if x2 <= x1:
+            x2 = min(w - 1, x1 + 1)
+        if y2 <= y1:
+            y2 = min(h - 1, y1 + 1)
+        return x1, y1, x2, y2
+
+    @staticmethod
+    def _area(bb: BoundingBox) -> int:
+        return max(0, bb.x2 - bb.x1) * max(0, bb.y2 - bb.y1)
+
+    @staticmethod
+    def _intersect_area(a: BoundingBox, b: BoundingBox) -> int:
+        ix1 = max(a.x1, b.x1)
+        iy1 = max(a.y1, b.y1)
+        ix2 = min(a.x2, b.x2)
+        iy2 = min(a.y2, b.y2)
+        if ix2 <= ix1 or iy2 <= iy1:
+            return 0
+        return (ix2 - ix1) * (iy2 - iy1)
+
+    @staticmethod
+    def _center(bb: BoundingBox) -> Tuple[float, float]:
+        return (0.5 * (bb.x1 + bb.x2), 0.5 * (bb.y1 + bb.y2))
+
+    @staticmethod
+    def _mean_hs(img_bgr: np.ndarray) -> Tuple[float, float]:
+        hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
+        return float(np.mean(hsv[:, :, 0])), float(np.mean(hsv[:, :, 1]))
+
+    def _hs_feature_from_roi(self, img_bgr: np.ndarray, box: BoundingBox) -> np.ndarray:
+        H, W = img_bgr.shape[:2]
+        x1, y1, x2, y2 = self._clip_box_to_image(box.x1, box.y1, box.x2, box.y2, W, H)
+        roi = img_bgr[y1:y2, x1:x2]
+        if roi.size == 0:
+            return np.array([0.0, 0.0], dtype=np.float32)
+        hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
+        lower_green = np.array([35, 60, 60], dtype=np.uint8)
+        upper_green = np.array([85, 255, 255], dtype=np.uint8)
+        green_mask = cv2.inRange(hsv, lower_green, upper_green)
+        non_green_mask = cv2.bitwise_not(green_mask)
+        num_non_green = int(np.count_nonzero(non_green_mask))
+        total = hsv.shape[0] * hsv.shape[1]
+        if num_non_green > max(50, total // 20):
+            h_vals = hsv[:, :, 0][non_green_mask > 0]
+            s_vals = hsv[:, :, 1][non_green_mask > 0]
+            h_mean = float(np.mean(h_vals)) if h_vals.size else 0.0
+            s_mean = float(np.mean(s_vals)) if s_vals.size else 0.0
+        else:
+            h_mean, s_mean = self._mean_hs(roi)
+        return np.array([h_mean, s_mean], dtype=np.float32)
+
+    def _ioa(self, a: BoundingBox, b: BoundingBox) -> float:
+        inter = self._intersect_area(a, b)
+        aa = self._area(a)
+        if aa <= 0:
+            return 0.0
+        return inter / aa
+
+    def suppress_quasi_total_containment(self, boxes: List[BoundingBox]) -> List[BoundingBox]:
+        if len(boxes) <= 1:
+            return boxes
+        keep = [True] * len(boxes)
+        for i in range(len(boxes)):
+            if not keep[i]:
+                continue
+            for j in range(len(boxes)):
+                if i == j or not keep[j]:
+                    continue
+                ioa_i_in_j = self._ioa(boxes[i], boxes[j])
+                if ioa_i_in_j >= self.QUASI_TOTAL_IOA:
+                    keep[i] = False
+                    break
+        return [bb for bb, k in zip(boxes, keep) if k]
+
+    def suppress_small_contained(self, boxes: List[BoundingBox]) -> List[BoundingBox]:
+        if len(boxes) <= 1:
+            return boxes
+        keep = [True] * len(boxes)
+        areas = [self._area(bb) for bb in boxes]
+        for i in range(len(boxes)):
+            if not keep[i]:
+                continue
+            for j in range(len(boxes)):
+                if i == j or not keep[j]:
+                    continue
+                ai, aj = areas[i], areas[j]
+                if ai == 0 or aj == 0:
+                    continue
+                if ai <= aj:
+                    ratio = ai / aj
+                    if ratio <= self.SMALL_RATIO_MAX:
+                        ioa_i_in_j = self._ioa(boxes[i], boxes[j])
+                        if ioa_i_in_j >= self.SMALL_CONTAINED_IOA:
+                            keep[i] = False
+                            break
+                else:
+                    ratio = aj / ai
+                    if ratio <= self.SMALL_RATIO_MAX:
+                        ioa_j_in_i = self._ioa(boxes[j], boxes[i])
+                        if ioa_j_in_i >= self.SMALL_CONTAINED_IOA:
+                            keep[j] = False
+        return [bb for bb, k in zip(boxes, keep) if k]
+
+    def _assign_players_two_clusters(self, features: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 20, 1.0)
+        _, labels, centers = cv2.kmeans(
+            np.float32(features),
+            K=2,
+            bestLabels=None,
+            criteria=criteria,
+            attempts=5,
+            flags=cv2.KMEANS_PP_CENTERS,
+        )
+        return labels.reshape(-1), centers
+
+    def _reclass_extra_goalkeepers(
+        self,
+        img_bgr: np.ndarray,
+        boxes: List[BoundingBox],
+        cluster_centers: Optional[np.ndarray],
+    ) -> None:
+        gk_idxs = [i for i, bb in enumerate(boxes) if int(bb.cls_id) == 1]
+        if len(gk_idxs) <= 1:
+            return
+        gk_idxs_sorted = sorted(gk_idxs, key=lambda i: boxes[i].conf, reverse=True)
+        keep_gk_idx = gk_idxs_sorted[0]
+        to_reclass = gk_idxs_sorted[1:]
+        for gki in to_reclass:
+            hs_gk = self._hs_feature_from_roi(img_bgr, boxes[gki])
+            if cluster_centers is not None:
+                d0 = float(np.linalg.norm(hs_gk - cluster_centers[0]))
+                d1 = float(np.linalg.norm(hs_gk - cluster_centers[1]))
+                assign_cls = 6 if d0 <= d1 else 7
+            else:
+                assign_cls = 6 if float(hs_gk[0]) < self.SINGLE_PLAYER_HUE_PIVOT else 7
+            boxes[gki].cls_id = int(assign_cls)
+
+    def predict_batch(
+        self,
+        batch_images: List[ndarray],
+        offset: int,
+        n_keypoints: int,
+        task_type: Optional[str] = None,
+    ) -> List[TVFrameResult]:
+        process_objects = task_type is None or task_type == "object"
+        process_keypoints = task_type is None or task_type == "keypoint"
+        bboxes: Dict[int, List[BoundingBox]] = {}
+        if process_objects:
+            bbox_model_results = self.bbox_model.predict(batch_images)
+        else:
+            bbox_model_results = None
+        if bbox_model_results is not None:
+            for frame_idx_in_batch, detection in enumerate(bbox_model_results):
+                if not hasattr(detection, "boxes") or detection.boxes is None:
+                    continue
+                boxes: List[BoundingBox] = []
+                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),
+                        )
+                    )
+                footballs = [bb for bb in boxes if int(bb.cls_id) == 0]
+                if len(footballs) > 1:
+                    best_ball = max(footballs, key=lambda b: b.conf)
+                    boxes = [bb for bb in boxes if int(bb.cls_id) != 0]
+                    boxes.append(best_ball)
+                boxes = self.suppress_quasi_total_containment(boxes)
+                boxes = self.suppress_small_contained(boxes)
+                img_bgr = batch_images[frame_idx_in_batch]
+                player_indices: List[int] = []
+                player_feats: List[np.ndarray] = []
+                for i, bb in enumerate(boxes):
+                    if int(bb.cls_id) == 2:
+                        hs = self._hs_feature_from_roi(img_bgr, bb)
+                        player_indices.append(i)
+                        player_feats.append(hs)
+                cluster_centers: Optional[np.ndarray] = None
+                n_players = len(player_feats)
+                if n_players >= 2:
+                    feats = np.vstack(player_feats)
+                    labels, centers = self._assign_players_two_clusters(feats)
+                    order = np.argsort(centers[:, 0])
+                    centers = centers[order]
+                    remap = {old_idx: new_idx for new_idx, old_idx in enumerate(order)}
+                    labels = np.vectorize(remap.get)(labels)
+                    cluster_centers = centers
+                    for idx_in_list, lbl in zip(player_indices, labels):
+                        boxes[idx_in_list].cls_id = 6 if int(lbl) == 0 else 7
+                elif n_players == 1:
+                    hue, _ = player_feats[0]
+                    boxes[player_indices[0]].cls_id = 6 if float(hue) < self.SINGLE_PLAYER_HUE_PIVOT else 7
+                self._reclass_extra_goalkeepers(img_bgr, boxes, cluster_centers)
+                bboxes[offset + frame_idx_in_batch] = boxes
+        keypoints: Dict[int, List[Tuple[int, int]]] = {}
+        if process_keypoints:
+            keypoints_model_results = self.keypoints_model.predict(batch_images)
+        else:
+            keypoints_model_results = None
+        if keypoints_model_results is not None:
+            for frame_idx_in_batch, detection in enumerate(keypoints_model_results):
+                if not hasattr(detection, "keypoints") or detection.keypoints is None:
+                    continue
+                frame_keypoints_with_conf: List[Tuple[int, int, float]] = []
+                for i, part_points in enumerate(detection.keypoints.data):
+                    for k_id, (x, y, _) in enumerate(part_points):
+                        confidence = float(detection.keypoints.conf[i][k_id])
+                        frame_keypoints_with_conf.append((int(x), int(y), confidence))
+                if len(frame_keypoints_with_conf) < n_keypoints:
+                    frame_keypoints_with_conf.extend(
+                        [(0, 0, 0.0)] * (n_keypoints - len(frame_keypoints_with_conf))
+                    )
+                else:
+                    frame_keypoints_with_conf = frame_keypoints_with_conf[:n_keypoints]
+                filtered_keypoints: List[Tuple[int, int]] = []
+                for idx, (x, y, confidence) in enumerate(frame_keypoints_with_conf):
+                    if idx in self.CORNER_INDICES:
+                        if confidence < 0.3:
+                            filtered_keypoints.append((0, 0))
+                        else:
+                            filtered_keypoints.append((int(x), int(y)))
+                    else:
+                        if confidence < 0.5:
+                            filtered_keypoints.append((0, 0))
+                        else:
+                            filtered_keypoints.append((int(x), int(y)))
+                keypoints[offset + frame_idx_in_batch] = filtered_keypoints
+        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)],
+                    ),
+                )
+            )
+        return results

objdetect.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8bbacfcb38e38b1b8816788e9e6e845160533719a0b87b693d58b932380d0d28
+size 152961687