futheros/render.py

"""Renderers for a FUT-HEROS analysis.

render()        -> clean analysis video: player segmentation masks (any number, bright on
                   possession), green goal mask, smooth ball trail, red goal-spot marker.
                   No HUD text, no pose card baked in (those are separate UI cards).
pose_card_img() -> the standalone Pose-Detection card (gray silhouette + black skeleton
                   on white), for display beside the video.
"""
from __future__ import annotations

import os
import shutil
import subprocess

import cv2
import numpy as np

from .pipeline import POSE_EDGES, L_ANKLE, R_ANKLE, AnalysisResult

# player palette (cycles for any number of players)
PAL = [(200, 90, 0), (0, 130, 255), (60, 180, 75), (180, 60, 200), (0, 200, 200), (120, 120, 0)]
PAL_HI = [(255, 170, 40), (40, 200, 255), (110, 240, 130), (230, 120, 250), (60, 255, 255), (200, 200, 60)]


def _mask_full(pmasks, key, f, W, H):
    m = pmasks if key == "goal" else (pmasks[key][f] if key in pmasks else None)
    if m is None or not m.any():
        return None
    return cv2.resize(m.astype(np.uint8), (W, H), interpolation=cv2.INTER_NEAREST).astype(bool)


def _blend(fr, mask, color, alpha):
    fr[mask] = (alpha * np.array(color, np.float32) + (1 - alpha) * fr[mask]).astype(np.uint8)


def render(clip: str, tracks: dict, res: AnalysisResult, out_path: str, truth: dict | None = None) -> str:
    W, H, n, fps = tracks["W"], tracks["H"], tracks["n"], tracks.get("fps") or 30.0
    xy = res.xy; holder = res.holder; pids = list(tracks["players"])
    gf = int(res.goal_time * fps) if res.goal_time is not None else None
    goal = tracks["goal"]; pmasks = tracks.get("pmasks", {}); gmask = tracks.get("goal_mask")
    # Stream annotated raw frames straight into ffmpeg via stdin. Avoids cv2.VideoWriter (doesn't
    # finalize off the main thread) AND per-frame JPEG files (slow IO) -> fast + thread-robust.
    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
    cap = cv2.VideoCapture(clip)
    ff = subprocess.Popen(
        ["ffmpeg", "-y", "-v", "error", "-f", "rawvideo", "-pix_fmt", "bgr24",
         "-s", f"{W}x{H}", "-r", f"{fps}", "-i", "-", "-an", "-c:v", "libx264",
         "-preset", "veryfast", "-crf", "23", "-pix_fmt", "yuv420p",
         "-movflags", "+faststart", out_path],
        stdin=subprocess.PIPE)
    tf = int(0.6 * fps); f = 0
    while True:
        ok, fr = cap.read()
        if not ok:
            break
        # goal segmentation (green)
        gm = _mask_full(gmask, "goal", f, W, H) if gmask is not None else None
        if gm is not None:
            _blend(fr, gm, (0, 200, 0), 0.40)
        elif goal is not None:
            cv2.rectangle(fr, (int(goal[0]), int(goal[1])), (int(goal[2]), int(goal[3])), (0, 200, 0), 2)
        # player segmentation masks (any number), brighter on possession, with P# labels
        for i, pid in enumerate(pids):
            has = holder[f] == pid
            col = PAL[i % len(PAL)]
            label = f"P{i+1}" + ("  BALL" if has else "")
            pmf = _mask_full(pmasks, pid, f, W, H)
            if pmf is not None:
                _blend(fr, pmf, PAL_HI[i % len(PAL_HI)] if has else col, 0.65 if has else 0.42)
                ys, xs = np.where(pmf)
                if len(xs):
                    lx, ly = int(np.median(xs)) - 14, max(18, int(ys.min()) - 8)
                    cv2.putText(fr, label, (lx, ly), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 3, cv2.LINE_AA)
                    cv2.putText(fr, label, (lx, ly), cv2.FONT_HERSHEY_SIMPLEX, 0.6, col, 2, cv2.LINE_AA)
            elif f in tracks["players"][pid]:
                cx, cy, w, h = tracks["players"][pid][f]
                cv2.rectangle(fr, (int(cx-w/2), int(cy-h/2)), (int(cx+w/2), int(cy+h/2)), col, 2)
                cv2.putText(fr, label, (int(cx-w/2), int(cy-h/2)-8), cv2.FONT_HERSHEY_SIMPLEX, 0.6, col, 2, cv2.LINE_AA)
        # ball trail
        pts = [(ff, xy[ff, 0], xy[ff, 1]) for ff in range(max(0, f-tf), f+1) if not np.isnan(xy[ff, 0])]
        for a, c in zip(pts, pts[1:]):
            cv2.line(fr, (int(a[1]), int(a[2])), (int(c[1]), int(c[2])), (0, 230, 255), 3, cv2.LINE_AA)
        # goal-spot marker (when scored)
        if res.goal == "goal" and gf is not None and goal is not None and f >= gf:
            gx, gy = (int(xy[gf, 0]), int(xy[gf, 1])) if not np.isnan(xy[gf, 0]) else (int((goal[0]+goal[2])/2), int((goal[1]+goal[3])/2))
            cv2.circle(fr, (gx, gy), 12, (0, 0, 255), 3, cv2.LINE_AA)
        if fr.shape[1] != W or fr.shape[0] != H:
            fr = cv2.resize(fr, (W, H))
        ff.stdin.write(fr.tobytes()); f += 1
    cap.release(); ff.stdin.close(); ff.wait()
    return out_path


def pose_card_img(clip: str, tracks: dict, res: AnalysisResult, out_png: str, card_h: int = 640) -> str | None:
    """Standalone 'Pose Detection' card: dark-grey silhouette (from the shooter's mask) +
    black MediaPipe skeleton on white. None if no pose was resolved."""
    if res.pose_pts is None or res.kbox is None or res.kick_frame is None:
        return None
    x0, y0, x1, y1 = res.kbox; w, h = x1 - x0, y1 - y0
    if w < 5 or h < 5:
        return None
    canvas = np.full((h, w, 3), 245, np.uint8)
    pm = tracks.get("pmasks", {}).get(res.shooter)
    if pm is not None:
        m = cv2.resize(pm[res.kick_frame].astype(np.uint8), (tracks["W"], tracks["H"]), cv2.INTER_NEAREST).astype(bool)
        sub = m[y0:y1, x0:x1]
        if sub.shape == (h, w):
            canvas[sub] = (90, 90, 90)
    lp = res.pose_pts - np.array([x0, y0])
    for a, b in POSE_EDGES:
        if a < len(lp) and b < len(lp):
            cv2.line(canvas, (int(lp[a][0]), int(lp[a][1])), (int(lp[b][0]), int(lp[b][1])), (25, 25, 25), 2, cv2.LINE_AA)
    for p in lp:
        cv2.circle(canvas, (int(p[0]), int(p[1])), 3, (25, 25, 25), -1)
    idx = L_ANKLE if res.foot == "L" else R_ANKLE if res.foot == "R" else None
    if idx is not None and idx < len(lp):
        cv2.circle(canvas, (int(lp[idx][0]), int(lp[idx][1])), 8, (0, 0, 255), 2, cv2.LINE_AA)
    sc = card_h / h
    card = cv2.resize(canvas, (max(1, int(w * sc)), card_h))
    os.makedirs(os.path.dirname(out_png) or ".", exist_ok=True)
    cv2.imwrite(out_png, card)
    return out_png