tests/test_robot_render.py

"""Tests for fpgm.robot.render / fpgm.robot.overlay.

Deliberately independent of :mod:`fpgm.robot.urdf` (which may not exist yet, and
in any case needs URDF/mesh assets this test suite must not depend on): the GL
tests below build a tiny synthetic one-link scene directly, exactly matching the
``link_meshes`` shape a real :class:`~fpgm.robot.urdf.RobotModel` would produce.
"""

from __future__ import annotations

from types import SimpleNamespace

import numpy as np
import pytest
import trimesh

from fpgm.geometry.camera import Camera
from fpgm.geometry.transforms import pose6_to_matrix
from fpgm.robot.overlay import composite, draw_link_dots, mask_to_bgr
from fpgm.robot.render import RenderError, RobotRenderer
from fpgm.types import CameraIntrinsics


def _build_renderer(
    link_meshes: dict[str, list[tuple[trimesh.Trimesh, np.ndarray]]], width: int, height: int
) -> RobotRenderer:
    """Construct a :class:`RobotRenderer`, skipping the test if OSMesa is unavailable."""
    try:
        return RobotRenderer(link_meshes, width, height)
    except RenderError as exc:  # pragma: no cover - environment-dependent
        pytest.skip(f"OSMesa offscreen GL context unavailable: {exc}")


class TestRobotRendererBoxDepth:
    """Reproduces the smoke test in the task brief: a unit box, straight-on camera."""

    def test_mask_and_depth_match_analytic_box(self) -> None:
        width, height = 320, 240
        fx, fy, cx, cy = 200.0, 200.0, 160.0, 120.0
        intrinsics = CameraIntrinsics(fx=fx, fy=fy, cx=cx, cy=cy, width=width, height=height)

        # Camera at world (0, 0, -4), identity rotation, looking down +Z (OpenCV
        # forward) toward the origin: world_to_cam = [I | t] with t = -camera_pos.
        world_to_cam = np.eye(4)
        world_to_cam[2, 3] = 4.0
        camera = Camera(intrinsics, world_to_cam)

        box = trimesh.creation.box(extents=[1.0, 1.0, 1.0])
        link_meshes = {"box": [(box, np.eye(4))]}
        link_poses = {"box": np.eye(4)}

        renderer = _build_renderer(link_meshes, width, height)
        try:
            result = renderer.render(link_poses, camera)
        finally:
            renderer.close()

        assert result.color.shape == (height, width, 3)
        assert result.depth.shape == (height, width)
        assert result.mask.shape == (height, width)
        assert result.mask.dtype == np.bool_

        # Front face of the unit box, viewed head-on from 4m: distance = 4 - 0.5 = 3.5m.
        # Projected half-width in pixels = fx * 0.5 / 3.5 ~= 28.6px -> a ~57x57 square.
        assert result.mask.any()
        assert 2800 <= int(result.mask.sum()) <= 3800
        assert result.depth[result.mask] == pytest.approx(3.5, abs=0.05)
        # Nothing outside the mask should report depth.
        assert np.all(result.depth[~result.mask] == 0.0)

    def test_close_is_idempotent_and_context_manager_closes(self) -> None:
        width, height = 64, 64
        intrinsics = CameraIntrinsics(fx=50, fy=50, cx=32, cy=32, width=width, height=height)
        world_to_cam = np.eye(4)
        world_to_cam[2, 3] = 4.0
        camera = Camera(intrinsics, world_to_cam)
        box = trimesh.creation.box(extents=[1.0, 1.0, 1.0])
        link_meshes = {"box": [(box, np.eye(4))]}

        renderer = _build_renderer(link_meshes, width, height)
        with renderer:
            renderer.render({"box": np.eye(4)}, camera)
        renderer.close()  # closing an already-closed renderer must not raise


class TestCameraConventionRoundTrip:
    """The single most valuable test: catches an axis-flip that mirrors the robot.

    A camera-frame point that is off-axis in *both* x and y is projected two ways:
    analytically via `Camera.project` (already covered by tests/test_geometry.py),
    and by rendering a tiny marker mesh there and reading back where pyrender
    actually put it. Getting the OpenCV->OpenGL axis flip wrong mirrors one or
    both image axes, which moves an off-center point to a clearly different
    pixel -- an on-axis point would not detect that, since its mirror image is
    itself.
    """

    def test_rendered_pixel_matches_camera_project(self) -> None:
        width, height = 320, 240
        fx, fy, cx, cy = 300.0, 300.0, 160.0, 120.0
        intrinsics = CameraIntrinsics(fx=fx, fy=fy, cx=cx, cy=cy, width=width, height=height)
        # A non-axis-aligned extrinsic (same shape as tests/test_geometry.py's
        # _make_camera) so a single sign error can't cancel out by symmetry.
        world_to_cam = pose6_to_matrix(
            np.array([0.3, -0.2, 0.5]), np.array([0.1, 0.2, -0.15])
        )
        camera = Camera(intrinsics, world_to_cam)

        # Off-axis in both x and y, comfortably in front of the camera.
        point_cam = np.array([[0.4, 0.25, 3.0]])
        point_world = camera.cam_to_world(point_cam)[0]
        expected_uv, expected_depth = camera.project(point_world.reshape(1, 3))
        assert expected_depth[0] > 0  # sanity: point is in front of the camera
        assert 0 < expected_uv[0, 0] < width
        assert 0 < expected_uv[0, 1] < height

        marker = trimesh.creation.box(extents=[0.05, 0.05, 0.05])
        mesh_to_link = np.eye(4)
        mesh_to_link[:3, 3] = point_world
        link_meshes = {"marker": [(marker, mesh_to_link)]}

        renderer = _build_renderer(link_meshes, width, height)
        try:
            result = renderer.render({"marker": np.eye(4)}, camera)
        finally:
            renderer.close()

        assert result.mask.any(), "marker did not render at all -- camera likely faces away"
        ys, xs = np.nonzero(result.mask)
        rendered_uv = np.array([xs.mean(), ys.mean()])

        assert rendered_uv == pytest.approx(expected_uv[0], abs=1.5)
        assert float(result.depth[result.mask].mean()) == pytest.approx(
            float(expected_depth[0]), abs=0.1
        )


class TestComposite:
    def test_blends_only_inside_mask_and_converts_rgb_to_bgr(self) -> None:
        frame_bgr = np.full((4, 4, 3), [10, 20, 30], dtype=np.uint8)
        original = frame_bgr.copy()
        mask = np.zeros((4, 4), dtype=bool)
        mask[1:3, 1:3] = True
        color_rgb = np.full((4, 4, 3), [200, 100, 50], dtype=np.uint8)  # R=200 G=100 B=50
        result = SimpleNamespace(
            color=color_rgb, depth=np.zeros((4, 4), dtype=np.float32), mask=mask
        )

        out = composite(frame_bgr, result, alpha=1.0)

        # alpha=1.0 fully replaces masked pixels with the render, RGB->BGR swapped.
        assert np.all(out[mask] == np.array([50, 100, 200], dtype=np.uint8))
        # Untouched outside the mask -- byte-identical to the original frame.
        assert np.array_equal(out[~mask], frame_bgr[~mask])
        # The input frame itself must not be mutated.
        assert np.array_equal(frame_bgr, original)

    def test_partial_alpha_blends_towards_render_color(self) -> None:
        frame_bgr = np.zeros((2, 2, 3), dtype=np.uint8)
        mask = np.ones((2, 2), dtype=bool)
        color_rgb = np.full((2, 2, 3), [100, 100, 100], dtype=np.uint8)
        result = SimpleNamespace(color=color_rgb, depth=np.zeros((2, 2)), mask=mask)

        out = composite(frame_bgr, result, alpha=0.5)
        assert np.all(out == 50)

    def test_empty_mask_returns_frame_unchanged(self) -> None:
        frame_bgr = np.full((3, 3, 3), 77, dtype=np.uint8)
        mask = np.zeros((3, 3), dtype=bool)
        result = SimpleNamespace(
            color=np.zeros((3, 3, 3), dtype=np.uint8), depth=np.zeros((3, 3)), mask=mask
        )
        out = composite(frame_bgr, result)
        assert np.array_equal(out, frame_bgr)

    def test_shape_mismatch_raises(self) -> None:
        frame_bgr = np.zeros((4, 4, 3), dtype=np.uint8)
        result = SimpleNamespace(
            color=np.zeros((2, 2, 3), dtype=np.uint8),
            depth=np.zeros((2, 2)),
            mask=np.zeros((2, 2), dtype=bool),
        )
        with pytest.raises(ValueError):
            composite(frame_bgr, result)


class TestMaskToBgr:
    def test_fills_true_pixels_with_color(self) -> None:
        mask = np.array([[True, False], [False, True]])
        out = mask_to_bgr(mask, color=(1, 2, 3))
        assert out.shape == (2, 2, 3)
        assert tuple(out[0, 0]) == (1, 2, 3)
        assert tuple(out[0, 1]) == (0, 0, 0)
        assert tuple(out[1, 0]) == (0, 0, 0)
        assert tuple(out[1, 1]) == (1, 2, 3)

    def test_default_color_is_white(self) -> None:
        mask = np.array([[True]])
        out = mask_to_bgr(mask)
        assert tuple(out[0, 0]) == (255, 255, 255)


class TestDrawLinkDots:
    def test_draws_a_dot_per_point_without_mutating_input(self) -> None:
        frame_bgr = np.zeros((50, 50, 3), dtype=np.uint8)
        original = frame_bgr.copy()
        uv = np.array([[10.0, 10.0], [40.0, 40.0]])

        out = draw_link_dots(frame_bgr, uv, names=["shoulder", "wrist"])

        assert out.shape == frame_bgr.shape
        assert np.array_equal(frame_bgr, original)
        assert out[10, 10].any()  # something was drawn near the first point
        assert out[40, 40].any()

    def test_mismatched_names_length_raises(self) -> None:
        frame_bgr = np.zeros((10, 10, 3), dtype=np.uint8)
        uv = np.array([[1.0, 1.0], [2.0, 2.0]])
        with pytest.raises(ValueError):
            draw_link_dots(frame_bgr, uv, names=["only_one"])