scorevision: push artifact

miner.py

@@ -11,6 +11,7 @@ from typing import Any
 
 import cv2
 import numpy as np
+import torch
 from numpy import ndarray
 from pydantic import BaseModel
 
@@ -30,6 +31,207 @@ class TVFrameResult(BaseModel):
     keypoints: list[tuple[int, int]]
 
 
+# ==========================
+# Football template keypoints (order matters!)
+# Copied from: scorevision/vlm_pipeline/domain_specific_schemas/football.py
+# ==========================
+FOOTBALL_KEYPOINTS: list[tuple[int, int]] = [
+    (5, 5),  # 1
+    (5, 140),  # 2
+    (5, 250),  # 3
+    (5, 430),  # 4
+    (5, 540),  # 5
+    (5, 675),  # 6
+    (55, 250),  # 7
+    (55, 430),  # 8
+    (110, 340),  # 9
+    (165, 140),  # 10
+    (165, 270),  # 11
+    (165, 410),  # 12
+    (165, 540),  # 13
+    (527, 5),  # 14
+    (527, 253),  # 15
+    (527, 433),  # 16
+    (527, 675),  # 17
+    (888, 140),  # 18
+    (888, 270),  # 19
+    (888, 410),  # 20
+    (888, 540),  # 21
+    (940, 340),  # 22
+    (998, 250),  # 23
+    (998, 430),  # 24
+    (1045, 5),  # 25
+    (1045, 140),  # 26
+    (1045, 250),  # 27
+    (1045, 430),  # 28
+    (1045, 540),  # 29
+    (1045, 675),  # 30
+    (435, 340),  # 31
+    (615, 340),  # 32
+]
+
+
+def _clamp_box(x1: int, y1: int, x2: int, y2: int, w: int, h: int) -> tuple[int, int, int, int]:
+    x1 = max(0, min(w - 1, x1))
+    y1 = max(0, min(h - 1, y1))
+    x2 = max(0, min(w - 1, x2))
+    y2 = max(0, min(h - 1, y2))
+    if x2 <= x1:
+        x2 = min(w - 1, x1 + 1)
+    if y2 <= y1:
+        y2 = min(h - 1, y1 + 1)
+    return x1, y1, x2, y2
+
+
+def _center_crop(box: tuple[int, int, int, int], frac: float = 0.55) -> tuple[int, int, int, int]:
+    """Take a smaller crop (helps focus on jersey color vs grass)."""
+    x1, y1, x2, y2 = box
+    cx = (x1 + x2) / 2
+    cy = (y1 + y2) / 2
+    w = (x2 - x1) * frac
+    h = (y2 - y1) * frac
+    nx1 = int(cx - w / 2)
+    nx2 = int(cx + w / 2)
+    ny1 = int(cy - h / 2)
+    ny2 = int(cy + h / 2)
+    return nx1, ny1, nx2, ny2
+
+
+def _mean_hsv(bgr: np.ndarray, box: tuple[int, int, int, int]) -> tuple[float, float, float]:
+    h, w = bgr.shape[:2]
+    x1, y1, x2, y2 = _clamp_box(*box, w, h)
+    crop = bgr[y1:y2, x1:x2]
+    if crop.size == 0:
+        return (0.0, 0.0, 0.0)
+    hsv = cv2.cvtColor(crop, cv2.COLOR_BGR2HSV)
+    mean = hsv.reshape(-1, 3).mean(axis=0)
+    return float(mean[0]), float(mean[1]), float(mean[2])
+
+
+def _hsv_dist(a: tuple[float, float, float], b: tuple[float, float, float]) -> float:
+    # hue wrap-around: treat hue as circular
+    dh = abs(a[0] - b[0])
+    dh = min(dh, 180 - dh)
+    ds = abs(a[1] - b[1])
+    dv = abs(a[2] - b[2])
+    return float(dh * 1.0 + ds * 0.25 + dv * 0.25)
+
+
+def _two_centroids(colors: list[tuple[float, float, float]]) -> tuple[tuple[float, float, float], tuple[float, float, float]] | None:
+    """Tiny k-means(k=2) for team jersey colors."""
+    if len(colors) < 2:
+        return None
+    pts = np.array(colors, dtype=np.float32)
+    # init: farthest-in-hue from first point
+    idx1 = 0
+    idx2 = int(np.argmax(np.abs(pts[:, 0] - pts[0, 0])))
+    c1 = pts[idx1].copy()
+    c2 = pts[idx2].copy()
+    for _ in range(8):
+        d1 = np.linalg.norm(pts - c1, axis=1)
+        d2 = np.linalg.norm(pts - c2, axis=1)
+        a1 = pts[d1 <= d2]
+        a2 = pts[d1 > d2]
+        if len(a1) > 0:
+            c1 = a1.mean(axis=0)
+        if len(a2) > 0:
+            c2 = a2.mean(axis=0)
+    return (float(c1[0]), float(c1[1]), float(c1[2])), (float(c2[0]), float(c2[1]), float(c2[2]))
+
+
+def _detect_pitch_lines_mask(bgr: np.ndarray) -> np.ndarray:
+    """Binary mask (0/255) for likely pitch lines."""
+    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
+    green = cv2.inRange(hsv, (25, 30, 30), (95, 255, 255))
+    green = cv2.medianBlur(green, 7)
+    white = cv2.inRange(hsv, (0, 0, 170), (180, 60, 255))
+    white = cv2.bitwise_and(white, white, mask=green)
+    k = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+    white = cv2.morphologyEx(white, cv2.MORPH_OPEN, k, iterations=1)
+    white = cv2.dilate(white, k, iterations=1)
+    return white
+
+
+def _load_template_line_mask(path_hf_repo: Path) -> tuple[np.ndarray, tuple[int, int]]:
+    """
+    Load football pitch template line mask from your miner repo folder.
+
+    You MUST copy `football_pitch_template.png` into `my_miner_repo/` before pushing.
+    """
+    tpl_name = os.getenv("PITCH_TEMPLATE_PNG", "football_pitch_template.png")
+    tpl_path = (path_hf_repo / tpl_name).resolve()
+    if not tpl_path.is_file():
+        raise FileNotFoundError(
+            f"Missing {tpl_name} in miner repo. Copy it from turbovision: "
+            f"scorevision/vlm_pipeline/domain_specific_schemas/football_pitch_template.png"
+        )
+    tpl = cv2.imread(str(tpl_path))
+    if tpl is None:
+        raise RuntimeError(f"Failed to read template image: {tpl_path}")
+    gray = cv2.cvtColor(tpl, cv2.COLOR_BGR2GRAY)
+    _, lines = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
+    return lines, (tpl.shape[1], tpl.shape[0])  # (W,H)
+
+
+def _estimate_homography_ecc(template_lines: np.ndarray, frame_lines: np.ndarray) -> np.ndarray | None:
+    """Estimate template->frame homography using ECC alignment on binary line masks."""
+    max_w = 960
+    fh, fw = frame_lines.shape[:2]
+    scale = min(1.0, max_w / float(fw)) if fw > 0 else 1.0
+
+    def _resize(img: np.ndarray) -> np.ndarray:
+        if scale >= 0.999:
+            return img
+        return cv2.resize(img, (int(img.shape[1] * scale), int(img.shape[0] * scale)), interpolation=cv2.INTER_AREA)
+
+    tpl = _resize(template_lines)
+    frm = _resize(frame_lines)
+    tpl_f = tpl.astype(np.float32) / 255.0
+    frm_f = frm.astype(np.float32) / 255.0
+
+    warp = np.eye(3, dtype=np.float32)
+    criteria = (
+        cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,
+        int(os.getenv("ECC_ITERS", "80")),
+        float(os.getenv("ECC_EPS", "1e-5")),
+    )
+    try:
+        # Find warp that best aligns template to frame (warp(template) ≈ frame)
+        cv2.findTransformECC(
+            frm_f,
+            tpl_f,
+            warp,
+            cv2.MOTION_HOMOGRAPHY,
+            criteria,
+            inputMask=None,
+            gaussFiltSize=3,
+        )
+        if scale < 0.999:
+            S = np.array([[1 / scale, 0, 0], [0, 1 / scale, 0], [0, 0, 1]], dtype=np.float32)
+            warp = S @ warp
+        return warp
+    except Exception:
+        return None
+
+
+def _project_keypoints(warp_tpl_to_frame: np.ndarray | None, frame_h: int, frame_w: int, n_keypoints: int) -> list[tuple[int, int]]:
+    if warp_tpl_to_frame is None:
+        return [(0, 0) for _ in range(n_keypoints)]
+    pts = np.array(FOOTBALL_KEYPOINTS[:n_keypoints], dtype=np.float32).reshape(1, -1, 2)
+    try:
+        out = cv2.perspectiveTransform(pts, warp_tpl_to_frame)[0]
+    except Exception:
+        return [(0, 0) for _ in range(n_keypoints)]
+    res: list[tuple[int, int]] = []
+    for x, y in out:
+        xi, yi = int(round(float(x))), int(round(float(y)))
+        if xi < 0 or yi < 0 or xi >= frame_w or yi >= frame_h:
+            res.append((0, 0))
+        else:
+            res.append((xi, yi))
+    return res
+
+
 class Miner:
     """
     Your miner engine.
@@ -71,7 +273,7 @@ class Miner:
                 "Add it to chute_config.yml (pip install segment-anything)."
             ) from e
 
-        device = "cuda" if os.getenv("CUDA_VISIBLE_DEVICES") else "cpu"
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         self.sam = sam_model_registry[model_type](checkpoint=str(ckpt_path))
         self.sam.to(device=device)
 
@@ -85,10 +287,13 @@ class Miner:
         )
 
         # ---------------- Keypoints ----------------
-        # Placeholder: output all zeros unless you add a keypoint detector.
-        self.enable_keypoints = os.getenv("ENABLE_KEYPOINTS", "0").lower() in ("1", "true", "yes")
-        self._kp_model: Any | None = None
-        # If you have a keypoint model, load it here from path_hf_repo.
+        # We'll estimate pitch keypoints via template line alignment (ECC).
+        # This is OpenCV-only (no extra model), but not perfect.
+        self.enable_keypoints = os.getenv("ENABLE_KEYPOINTS", "1").lower() in ("1", "true", "yes")
+        self._template_lines: np.ndarray | None = None
+        self._template_wh: tuple[int, int] | None = None
+        if self.enable_keypoints:
+            self._template_lines, self._template_wh = _load_template_line_mask(path_hf_repo)
 
     def __repr__(self) -> str:
         return (
@@ -102,32 +307,43 @@ class Miner:
         offset: int,
         n_keypoints: int,
     ) -> list[TVFrameResult]:
-        # ------------------ Boxes (SAM masks -> boxes) ------------------
-        # SAM returns masks for "things" but does not label them (player/ref/ball).
-        # For a first working miner, we mark everything as "player" (cls_id=2).
-        # To score well, you will later need classification (ball/ref/goalkeeper/team).
         bboxes: dict[int, list[BoundingBox]] = {}
+        keypoints: dict[int, list[tuple[int, int]]] = {}
 
+        # Per-frame processing (keeps logic simple; you can optimize later)
         for i, img in enumerate(batch_images):
             frame_id = offset + i
 
-            # Convert BGR(OpenCV) -> RGB(SAM)
             if img is None:
                 bboxes[frame_id] = []
+                keypoints[frame_id] = [(0, 0) for _ in range(n_keypoints)]
                 continue
-            rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
+            frame_h, frame_w = img.shape[:2]
+
+            # ---------------- Keypoints (ECC template alignment) ----------------
+            if self.enable_keypoints and self._template_lines is not None:
+                frame_lines = _detect_pitch_lines_mask(img)
+                warp = _estimate_homography_ecc(self._template_lines, frame_lines)
+                keypoints[frame_id] = _project_keypoints(warp, frame_h, frame_w, n_keypoints)
+            else:
+                keypoints[frame_id] = [(0, 0) for _ in range(n_keypoints)]
+
+            # ---------------- Boxes (SAM masks -> boxes) ----------------
+            # SAM returns masks without semantic classes. We'll add simple heuristics:
+            # - ball: very small near-square bbox (cls_id=0)
+            # - teams: cluster jersey colors into two groups -> cls_id=6 or 7
+            # - non-team dark/odd colors -> referee (cls_id=3) and maybe one goalkeeper (cls_id=1)
+            rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
             masks = self.mask_generator.generate(rgb)  # list[dict]
 
-            # Filter out giant masks (often the grass/background) and tiny noise.
-            H, W = rgb.shape[:2]
-            area_frame = float(H * W)
-            out_boxes: list[BoundingBox] = []
+            area_frame = float(frame_h * frame_w)
+            cand: list[tuple[int, int, int, int, float]] = []
+
             for m in masks:
-                # segment-anything returns bbox as [x, y, w, h]
-                x, y, w, h = m.get("bbox") or (0, 0, 0, 0)
+                x, y, w_, h_ = m.get("bbox") or (0, 0, 0, 0)
                 x1, y1 = int(x), int(y)
-                x2, y2 = int(x + w), int(y + h)
+                x2, y2 = int(x + w_), int(y + h_)
                 if x2 <= x1 or y2 <= y1:
                     continue
 
@@ -137,26 +353,94 @@ class Miner:
                 if box_area / area_frame > float(os.getenv("MAX_BOX_AREA_FRAC", "0.25")):
                     continue
 
+                # human-ish aspect ratio filter (helps remove lines/signage)
+                ar = float((y2 - y1) / max(1.0, (x2 - x1)))
+                if ar < float(os.getenv("MIN_ASPECT_RATIO", "1.0")):
+                    continue
+                if ar > float(os.getenv("MAX_ASPECT_RATIO", "6.0")):
+                    continue
+
                 conf = float(m.get("predicted_iou") or 0.5)
-                out_boxes.append(
-                    BoundingBox(
-                        x1=x1,
-                        y1=y1,
-                        x2=x2,
-                        y2=y2,
-                        cls_id=2,  # default: player
-                        conf=conf,
-                    )
-                )
+                cand.append((x1, y1, x2, y2, conf))
 
-            bboxes[frame_id] = out_boxes
+            # ball candidates (tiny)
+            ball_max_area = int(os.getenv("BALL_MAX_AREA", "900"))
+            ball: list[BoundingBox] = []
+            people: list[tuple[int, int, int, int, float]] = []
+            for x1, y1, x2, y2, conf in cand:
+                bw = x2 - x1
+                bh = y2 - y1
+                a = bw * bh
+                if a <= ball_max_area and 0.6 <= (bw / max(1.0, bh)) <= 1.6:
+                    ball.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, cls_id=0, conf=float(conf)))
+                else:
+                    people.append((x1, y1, x2, y2, conf))
 
-        # ---------------- Keypoints (length = n_keypoints) ----------------
-        keypoints: dict[int, list[tuple[int, int]]] = {}
-        # Placeholder (zeros). Replace with your own keypoint detector when ready.
-        for i in range(len(batch_images)):
-            frame_id = offset + i
-            keypoints[frame_id] = [(0, 0) for _ in range(n_keypoints)]
+            # keep 1 ball max
+            ball.sort(key=lambda b: b.conf, reverse=True)
+            if len(ball) > 1:
+                ball = ball[:1]
+
+            # jersey colors (center crop per person box)
+            colors: list[tuple[float, float, float]] = []
+            for x1, y1, x2, y2, _conf in people:
+                cx1, cy1, cx2, cy2 = _center_crop((x1, y1, x2, y2))
+                cx1, cy1, cx2, cy2 = _clamp_box(cx1, cy1, cx2, cy2, frame_w, frame_h)
+                colors.append(_mean_hsv(img, (cx1, cy1, cx2, cy2)))
+
+            cents = _two_centroids(colors)
+            team1_c, team2_c = cents if cents is not None else ((0.0, 0.0, 0.0), (90.0, 0.0, 0.0))
+
+            dark_v_thresh = float(os.getenv("REF_DARK_V", "70"))
+            nonteam_dist = float(os.getenv("NONTEAM_DIST", "45"))
+
+            team_boxes: list[BoundingBox] = []
+            nonteam_boxes: list[tuple[int, int, int, int, float, tuple[float, float, float]]] = []
+
+            for (x1, y1, x2, y2, conf), c in zip(people, colors, strict=False):
+                d1 = _hsv_dist(c, team1_c)
+                d2 = _hsv_dist(c, team2_c)
+                if c[2] < dark_v_thresh or (d1 > nonteam_dist and d2 > nonteam_dist):
+                    nonteam_boxes.append((x1, y1, x2, y2, conf, c))
+                else:
+                    cls = 6 if d1 <= d2 else 7
+                    team_boxes.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, cls_id=int(cls), conf=float(conf)))
+
+            # goalkeeper vs referee split among nonteam:
+            # choose at most 1 goalkeeper near left/right third; rest referees.
+            gk_box: BoundingBox | None = None
+            refs: list[BoundingBox] = []
+            if nonteam_boxes:
+                edge_candidates = []
+                mid_candidates = []
+                for x1, y1, x2, y2, conf, _c in nonteam_boxes:
+                    cx = (x1 + x2) / 2.0
+                    if cx < frame_w * 0.33 or cx > frame_w * 0.66:
+                        edge_candidates.append((x1, y1, x2, y2, conf))
+                    else:
+                        mid_candidates.append((x1, y1, x2, y2, conf))
+                edge_candidates.sort(key=lambda t: t[4], reverse=True)
+                if edge_candidates:
+                    x1, y1, x2, y2, conf = edge_candidates[0]
+                    gk_box = BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, cls_id=1, conf=float(conf))
+                    for x1, y1, x2, y2, conf in edge_candidates[1:]:
+                        refs.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, cls_id=3, conf=float(conf)))
+                for x1, y1, x2, y2, conf in mid_candidates:
+                    refs.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, cls_id=3, conf=float(conf)))
+
+            out: list[BoundingBox] = []
+            out.extend(ball)
+            if gk_box is not None:
+                out.append(gk_box)
+            out.extend(team_boxes)
+            out.extend(refs)
+
+            max_boxes = int(os.getenv("MAX_BOXES_PER_FRAME", "40"))
+            if len(out) > max_boxes:
+                out.sort(key=lambda b: b.conf, reverse=True)
+                out = out[:max_boxes]
+
+            bboxes[frame_id] = out
 
         # ---------------- Combine ------------------
         results: list[TVFrameResult] = []