src/fpgm/robot/overlay.py

"""Composite a rendered robot onto video frames.

Pure numpy + OpenCV, deliberately with **no** ``pyrender`` import anywhere in this
module (even for typing -- ``RenderResult`` is only imported under
``TYPE_CHECKING``), so it stays importable and unit-testable on a box with no GL
context at all.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import cv2
import numpy as np

from fpgm.viz.overlays import color_for

if TYPE_CHECKING:
    from fpgm.robot.render import RenderResult


def composite(frame_bgr: np.ndarray, result: RenderResult, alpha: float = 0.85) -> np.ndarray:
    """Alpha-blend a rendered robot over a video frame, restricted to its mask.

    Args:
        frame_bgr: ``(H, W, 3)`` uint8 BGR video frame.
        result: Anything shaped like :class:`~fpgm.robot.render.RenderResult`,
            i.e. exposing ``.color`` (``(H, W, 3)`` uint8, RGB) and ``.mask``
            (``(H, W)`` bool). Only duck-typed on this side of the boundary --
            see the module docstring for why.
        alpha: Blend weight given to the rendered robot inside the mask; ``1.0``
            replaces the frame outright there, ``0.0`` leaves it untouched.

    Returns:
        ``(H, W, 3)`` uint8 BGR frame. Pixels outside ``result.mask`` are
        byte-identical to ``frame_bgr``.
    """
    if frame_bgr.shape[:2] != result.mask.shape:
        raise ValueError(
            f"frame shape {frame_bgr.shape[:2]} does not match mask shape "
            f"{result.mask.shape}"
        )
    out = frame_bgr.copy()
    mask = result.mask
    if not mask.any():
        return out

    render_bgr = np.ascontiguousarray(result.color[..., ::-1])  # pyrender is RGB, frames are BGR
    blended = cv2.addWeighted(frame_bgr, 1.0 - alpha, render_bgr, alpha, 0.0)
    out[mask] = blended[mask]
    return out


def mask_to_bgr(mask: np.ndarray, color: tuple[int, int, int] = (255, 255, 255)) -> np.ndarray:
    """Render a boolean mask as a 3-channel BGR image, for writing a mask video.

    Args:
        mask: ``(H, W)`` bool.
        color: BGR fill colour for ``True`` pixels; the rest stays black.
    """
    out = np.zeros((*mask.shape, 3), dtype=np.uint8)
    out[mask] = color
    return out


def draw_link_dots(
    frame_bgr: np.ndarray,
    uv: np.ndarray,
    names: list[str] | None = None,
    radius: int = 4,
) -> np.ndarray:
    """Draw one dot per link origin -- the cheap FK-only sanity check.

    Args:
        frame_bgr: ``(H, W, 3)`` uint8 BGR frame.
        uv: ``(N, 2)`` pixel coordinates, one row per link origin.
        names: Optional length-``N`` labels drawn next to each dot.
        radius: Dot radius in pixels.

    Returns:
        A copy of ``frame_bgr`` with the dots (and optional labels) drawn on.
    """
    out = frame_bgr.copy()
    uv = np.asarray(uv)
    if names is not None and len(names) != uv.shape[0]:
        raise ValueError(f"got {len(names)} names for {uv.shape[0]} points")

    for i in range(uv.shape[0]):
        color = color_for(i)
        center = tuple(np.round(uv[i]).astype(int))
        cv2.circle(out, center, radius, color, -1, lineType=cv2.LINE_AA)
        if names is not None:
            cv2.putText(
                out,
                names[i],
                (center[0] + radius + 2, center[1]),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.4,
                color,
                1,
                lineType=cv2.LINE_AA,
            )
    return out