miner.py

from pathlib import Path

# NOTE:
# - This is copied from `example_miner/miner.py` as a starting point.
# - This version shows how to use a SAM-style segmentation model as your detector.
# - SAM gives masks (segmentation). This subnet expects boxes, so we convert masks -> boxes.
# - SAM does NOT give 32 pitch keypoints; you likely need a separate keypoint model.

import os
from typing import Any

import cv2
import numpy as np
import torch
from numpy import ndarray
from pydantic import BaseModel


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]]


# ==========================
# 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.

    Requirements (must keep):
    - file name: `miner.py` (repo root)
    - class name: `Miner`
    - method: `predict_batch(batch_images, offset, n_keypoints)`
    """

    def __init__(self, path_hf_repo: Path) -> None:
        """
        Load your models from the HuggingFace repo snapshot directory.

        For SAM-based detection:
        - Put your SAM checkpoint file in this repo folder (same folder as miner.py)
        - Set SAM_CHECKPOINT env var (optional) to choose the filename.
        """
        self.path_hf_repo = path_hf_repo

        # ---------------- SAM setup ----------------
        # IMPORTANT: "SAM 3" can mean different things. This skeleton uses the common
        # Segment Anything API shape (sam_model_registry + SamAutomaticMaskGenerator).
        # If your SAM3 is different, keep the structure and replace the loading/inference.
        ckpt_name = os.getenv("SAM_CHECKPOINT", "sam_vit_h_4b8939.pth")
        ckpt_path = (path_hf_repo / ckpt_name).resolve()
        if not ckpt_path.is_file():
            raise FileNotFoundError(
                f"SAM checkpoint not found: {ckpt_path}. Put the checkpoint in your HF repo "
                f"and/or set SAM_CHECKPOINT to the correct filename."
            )

        model_type = os.getenv("SAM_MODEL_TYPE", "vit_h")  # vit_h / vit_l / vit_b (depends on checkpoint)
        try:
            from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
        except Exception as e:
            raise ImportError(
                "segment-anything is not installed in the Chutes image. "
                "Add it to chute_config.yml (pip install segment-anything)."
            ) from e

        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)

        # Tunables: lower points_per_side => faster, fewer masks.
        self.mask_generator = SamAutomaticMaskGenerator(
            self.sam,
            points_per_side=int(os.getenv("SAM_POINTS_PER_SIDE", "16")),
            pred_iou_thresh=float(os.getenv("SAM_PRED_IOU_THRESH", "0.88")),
            stability_score_thresh=float(os.getenv("SAM_STABILITY_THRESH", "0.90")),
            min_mask_region_area=int(os.getenv("SAM_MIN_REGION_AREA", "200")),
        )

        # ---------------- Keypoints ----------------
        # 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 (
            f"SAM: {type(self.sam).__name__}\n"
            f"Keypoints enabled: {self.enable_keypoints}"
        )

    def predict_batch(
        self,
        batch_images: list[ndarray],
        offset: int,
        n_keypoints: int,
    ) -> list[TVFrameResult]:
        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

            if img is None:
                bboxes[frame_id] = []
                keypoints[frame_id] = [(0, 0) for _ in range(n_keypoints)]
                continue

            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]

            area_frame = float(frame_h * frame_w)
            cand: list[tuple[int, int, int, int, float]] = []

            for m in masks:
                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_)
                if x2 <= x1 or y2 <= y1:
                    continue

                box_area = float((x2 - x1) * (y2 - y1))
                if box_area < float(os.getenv("MIN_BOX_AREA", "250")):
                    continue
                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)
                cand.append((x1, y1, x2, y2, conf))

            # 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))

            # 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] = []
        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