src/fpgm/robot/render.py

"""Offscreen rendering of the articulated robot mesh via ``pyrender`` + OSMesa.

This machine's NVIDIA driver has no hardware OpenGL available (kernel driver
present, but no ``libEGL_nvidia``/``libGLX_nvidia`` at a matching version), so
every render here goes through Mesa's software rasteriser (OSMesa) instead of
EGL. That requires two workarounds, both applied by :func:`_configure_headless_gl`
*before* ``pyrender`` (or ``OpenGL``) is imported anywhere in the process --
see that function's docstring for why the ordering matters.

``RobotRenderer`` builds its ``pyrender`` scene graph once and only updates node
poses per frame (never rebuilds it), because scene construction dominates
runtime on software GL.
"""

from __future__ import annotations

import os
import sys
import types
from dataclasses import dataclass

import numpy as np
import trimesh

from fpgm.geometry.camera import Camera
from fpgm.types import FpgmError
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)


class RenderError(FpgmError):
    """Raised when the OSMesa offscreen GL context cannot be created."""


def _configure_headless_gl() -> None:
    """Make ``pyrender`` importable and offscreen-renderable on a headless box.

    Two independent workarounds, both required, both order-sensitive:

    1. ``PYOPENGL_PLATFORM=osmesa`` must be set **before** ``pyrender``/``OpenGL``
       is first imported anywhere in the process -- ``PyOpenGL`` reads this
       environment variable at import time to decide which GL binding to load,
       and it cannot be changed afterwards. ``setdefault`` is used so a caller
       that has already chosen a different platform (e.g. ``egl`` on a machine
       that does have a working driver) is not overridden.
    2. ``pyrender/__init__.py`` unconditionally does ``from .viewer import
       Viewer``, which imports ``pyglet`` -> ``pyglet.window`` -> ``libXrender``
       for the *interactive* viewer window -- and dies with an ``AttributeError``
       on a box with no X server, even though this module never uses that
       viewer (only :class:`pyrender.OffscreenRenderer`). Pre-registering a
       stub ``pyrender.viewer`` module short-circuits that import before it can
       run, without patching ``pyrender`` itself.
    """
    os.environ.setdefault("PYOPENGL_PLATFORM", "osmesa")
    if "pyrender.viewer" not in sys.modules:
        stub = types.ModuleType("pyrender.viewer")
        stub.Viewer = None  # type: ignore[attr-defined]
        sys.modules["pyrender.viewer"] = stub


_configure_headless_gl()

import pyrender  # noqa: E402  (must follow _configure_headless_gl(), see above)

#: pyrender/OpenGL cameras look down -Z with +Y up; our extrinsics are OpenCV
#: (+Z-forward, +Y-down). Right-multiplying a cam-to-world pose by this flips the
#: Y and Z camera axes to go from one convention to the other.
_GL_AXIS_FLIP = np.diag([1.0, -1.0, -1.0, 1.0])

#: Origin plus the three unit basis points, used to recover an affine matrix from
#: a point-transform function -- see :func:`_camera_to_world_matrix`.
_BASIS_POINTS = np.array(
    [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
)


@dataclass
class RenderResult:
    """One rendered frame of the robot."""

    color: np.ndarray  # (H, W, 3) uint8, RGB
    depth: np.ndarray  # (H, W) float32, metres; 0 where no geometry was hit
    mask: np.ndarray  # (H, W) bool


def _camera_to_world_matrix(camera: Camera) -> np.ndarray:
    """Reconstruct `camera`'s ``(4, 4)`` camera-to-world SE3 matrix.

    :class:`~fpgm.geometry.camera.Camera` deliberately exposes only point-transform
    methods, not its raw extrinsic matrix. The matrix is recovered here by
    transforming the origin and the three axis basis points through
    :meth:`Camera.cam_to_world` -- exact for any affine map -- rather than reaching
    into ``Camera``'s private attributes.
    """
    transformed = camera.cam_to_world(_BASIS_POINTS)  # (4, 3)
    translation = transformed[0]
    rotation = (transformed[1:] - translation).T
    matrix = np.eye(4, dtype=np.float64)
    matrix[:3, :3] = rotation
    matrix[:3, 3] = translation
    return matrix


def _gl_camera_pose(camera: Camera) -> np.ndarray:
    """Convert `camera`'s OpenCV world-to-camera extrinsic to a pyrender cam2world pose."""
    return _camera_to_world_matrix(camera) @ _GL_AXIS_FLIP


class RobotRenderer:
    """Offscreen software-GL renderer for a rigged, multi-link robot mesh.

    The ``pyrender`` scene graph (one node per visual geometry, plus one camera
    node and one light node) is built once in :meth:`__init__`. Each call to
    :meth:`render` only updates node poses via ``scene.set_pose`` -- rebuilding
    the scene every frame would dominate runtime on OSMesa's software rasteriser.
    """

    def __init__(
        self,
        link_meshes: dict[str, list[tuple[trimesh.Trimesh, np.ndarray]]],
        width: int,
        height: int,
        *,
        ambient: float = 0.4,
        bg_color: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
        znear: float = 0.01,
        zfar: float = 10.0,
        light_intensity: float = 3.0,
    ) -> None:
        """Build the persistent scene graph.

        Args:
            link_meshes: link name -> list of ``(mesh, mesh_to_link)`` pairs, where
                ``mesh_to_link`` is the ``(4, 4)`` transform from the mesh's own
                vertex frame into its link's frame. Injected rather than loaded
                from a URDF here, so this class stays testable with a synthetic
                one-link scene and independent of :mod:`fpgm.robot.urdf`.
            width: Render width in pixels.
            height: Render height in pixels.
            ambient: Flat ambient light term applied everywhere (0-1), so the
                unlit side of the robot is not pure black.
            bg_color: RGBA background colour. Irrelevant to the mask (which is
                derived from depth, not colour/alpha -- see :meth:`render`), but
                affects the ``color`` channel outside the robot.
            znear: Near clip plane, metres.
            zfar: Far clip plane, metres.
            light_intensity: Intensity of the directional light attached at the
                camera pose (see :meth:`render`); without it the render is black.

        Raises:
            RenderError: if the OSMesa offscreen GL context cannot be created.
        """
        self.width = int(width)
        self.height = int(height)
        self._znear = float(znear)
        self._zfar = float(zfar)

        self.scene = pyrender.Scene(
            ambient_light=np.full(3, float(ambient), dtype=np.float64),
            bg_color=list(bg_color),
        )

        # One node per visual geometry, built once; (node, mesh_to_link) is kept
        # per link so render() only ever computes `link_pose @ mesh_to_link`.
        self._link_nodes: dict[str, list[tuple[pyrender.Node, np.ndarray]]] = {}
        for link, meshes in link_meshes.items():
            entries: list[tuple[pyrender.Node, np.ndarray]] = []
            for mesh, mesh_to_link in meshes:
                mesh_to_link = np.asarray(mesh_to_link, dtype=np.float64)
                # smooth=False: some of this URDF's COLLADA visuals carry per-face
                # (not per-vertex) colours, and pyrender.Mesh.from_trimesh
                # unconditionally raises ValueError('Cannot use face colors with a
                # smooth mesh') for those unless smooth=False.
                pmesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
                node = pyrender.Node(mesh=pmesh, matrix=mesh_to_link)
                self.scene.add_node(node)
                entries.append((node, mesh_to_link))
            self._link_nodes[link] = entries

        placeholder_camera = pyrender.IntrinsicsCamera(
            fx=1.0, fy=1.0, cx=self.width / 2.0, cy=self.height / 2.0,
            znear=self._znear, zfar=self._zfar,
        )
        self._camera_node = pyrender.Node(camera=placeholder_camera, matrix=np.eye(4))
        self.scene.add_node(self._camera_node)

        # Lit from the viewpoint: attached at the camera pose every frame in
        # render(), so the visible side of the robot is never in shadow.
        light = pyrender.DirectionalLight(color=np.ones(3), intensity=float(light_intensity))
        self._light_node = pyrender.Node(light=light, matrix=np.eye(4))
        self.scene.add_node(self._light_node)

        self._renderer: pyrender.OffscreenRenderer | None = None
        try:
            self._renderer = pyrender.OffscreenRenderer(self.width, self.height)
        except Exception as exc:  # pragma: no cover - depends on the host's GL stack
            raise RenderError(
                "could not create an OSMesa offscreen GL context. RobotRenderer "
                "requires PYOPENGL_PLATFORM=osmesa (set by "
                "fpgm.robot.render._configure_headless_gl before pyrender is "
                "imported) plus a software-GL-capable PyOpenGL/OSMesa install."
            ) from exc

    def render(self, link_poses: dict[str, np.ndarray], camera: Camera) -> RenderResult:
        """Render the robot posed by `link_poses`, as seen by `camera`.

        Args:
            link_poses: link name -> ``(4, 4)`` pose in the base/world frame. Must
                contain an entry for every link this renderer was built with
                geometry for.
            camera: Camera to render from; its intrinsics' resolution must match
                the resolution this renderer was constructed with.

        Returns:
            A :class:`RenderResult` with the mask derived from depth (see below),
            never from an alpha channel.

        Raises:
            ValueError: if `camera`'s resolution does not match this renderer's.
            KeyError: if `link_poses` is missing a link this renderer has geometry
                for.
        """
        if camera.K.width != self.width or camera.K.height != self.height:
            raise ValueError(
                f"camera resolution ({camera.K.width}x{camera.K.height}) does not "
                f"match renderer resolution ({self.width}x{self.height})"
            )

        for link, entries in self._link_nodes.items():
            try:
                link_pose = np.asarray(link_poses[link], dtype=np.float64)
            except KeyError:
                raise KeyError(f"link_poses is missing a pose for link {link!r}") from None
            for node, mesh_to_link in entries:
                self.scene.set_pose(node, link_pose @ mesh_to_link)

        gl_pose = _gl_camera_pose(camera)
        # Intrinsics are swapped in wholesale (rather than mutated in place) since
        # pyrender.IntrinsicsCamera is otherwise immutable.
        self._camera_node.camera = pyrender.IntrinsicsCamera(
            fx=camera.K.fx, fy=camera.K.fy, cx=camera.K.cx, cy=camera.K.cy,
            znear=self._znear, zfar=self._zfar,
        )
        self.scene.set_pose(self._camera_node, gl_pose)
        self.scene.set_pose(self._light_node, gl_pose)

        color, depth = self._renderer.render(self.scene)
        # Measured on this OSMesa build: passing RenderFlags.RGBA still returns a
        # 3-channel array, not 4 -- so the mask is derived from depth (reliable),
        # never from an alpha channel (silently absent).
        mask = depth > 0.0
        return RenderResult(
            color=np.asarray(color, dtype=np.uint8),
            depth=np.asarray(depth, dtype=np.float32),
            mask=mask,
        )

    def close(self) -> None:
        """Release the underlying GL context. Safe to call more than once."""
        if self._renderer is not None:
            self._renderer.delete()
            self._renderer = None

    def __enter__(self) -> RobotRenderer:
        return self

    def __exit__(self, *exc_info: object) -> None:
        self.close()