scripts/render_robot.py

#!/usr/bin/env python
"""Render the articulated robot mesh over a DROID episode's exterior camera video.

    PYTHONPATH=src python scripts/render_robot.py --episode <uuid> --camera ext1

Reads joint/gripper state and per-frame camera extrinsics straight from a local
``data/droid_raw/<uuid>/trajectory.h5`` (no download attempted -- these assets are
expected to already be on disk), forward-kinematics them through a
:class:`~fpgm.robot.urdf.RobotModel`, and renders the posed mesh with
:class:`~fpgm.robot.render.RobotRenderer` (OSMesa offscreen GL; see that module's
docstring for why hardware GL is not used on this box).

Two modes:

* Full render (default): writes ``robot_overlay.mp4`` (video composited with the
  rendered robot), ``robot_mask.mp4`` (mask visualised as white-on-black), and
  ``robot_mask.npz`` (the same masks, bit-packed).
* ``--fk-only``: skips mesh loading and rendering entirely and instead projects
  each link's origin into the image as a coloured dot (``robot_fk_dots.mp4``) --
  a cheap way to sanity-check the kinematic chain and camera convention before
  paying for full rendering.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import h5py
import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from fpgm.data.droid_raw import cam2world_vector_to_world2cam, read_mp4_properties  # noqa: E402
from fpgm.data.pointworld import FlowsReader  # noqa: E402
from fpgm.geometry.camera import Camera  # noqa: E402
from fpgm.robot.overlay import composite, draw_link_dots, mask_to_bgr  # noqa: E402
from fpgm.types import CameraIntrinsics, DataError  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402
from fpgm.viz.overlays import VideoWriter, read_frames_bgr  # noqa: E402

logger = get_logger("render_robot")

#: Falls back to this annotation resolution when a clip has no decoded
#: `initial_rgb` to measure it from -- see `_read_base_intrinsics`.
_FALLBACK_ANNOTATION_SIZE = (320, 180)

_DEFAULT_URDF = (
    REPO_ROOT
    / "third_party"
    / "robot_description"
    / "pointworld_franka_robotiq_2f85"
    / "franka_panda_robotiq_2f85.urdf"
)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--episode", required=True, help="DROID episode uuid")
    p.add_argument(
        "--camera", default="ext1", choices=["ext1", "ext2"], help="exterior camera role"
    )
    p.add_argument(
        "--max-frames", type=int, default=None, help="cap the number of frames processed"
    )
    p.add_argument(
        "--urdf",
        type=Path,
        default=_DEFAULT_URDF,
        help="robot URDF path (default: PointWorld's DROID Franka+Robotiq asset)",
    )
    p.add_argument("--stride", type=int, default=1, help="process every Nth trajectory step")
    p.add_argument(
        "--outputs",
        type=Path,
        default=None,
        help="output directory (default: outputs/<episode-uuid>/robot)",
    )
    p.add_argument(
        "--fk-only",
        action="store_true",
        help="skip mesh rendering; project link origins as dots (cheap sanity check)",
    )
    return p.parse_args()


def _episode_dir(uuid: str) -> Path:
    return REPO_ROOT / "data" / "droid_raw" / uuid


def _flows_path(uuid: str) -> Path:
    return REPO_ROOT / "data" / "pointworld" / "droid" / "flows-fs-optimized" / f"{uuid}_flows.h5"


def _load_metadata(uuid: str) -> dict:
    path = _episode_dir(uuid) / "metadata.json"
    if not path.exists():
        raise DataError(
            f"no local DROID metadata for episode {uuid!r} at {path}; this script "
            "does not download -- fetch the episode first."
        )
    return json.loads(path.read_text())


def _read_base_intrinsics(uuid: str, camera_serial: str) -> CameraIntrinsics:
    """Read the PointWorld scene-flow intrinsic for `camera_serial`, at its native resolution.

    Raises:
        DataError: if the episode's flows.h5 is absent, or has no clip for
            `camera_serial` -- this intrinsic is the only source of calibration
            this script has, so there is no fallback.
    """
    flows_path = _flows_path(uuid)
    if not flows_path.exists():
        raise DataError(
            f"PointWorld flows.h5 not found for episode {uuid!r} at {flows_path}; "
            "cannot recover camera intrinsics without it."
        )
    with FlowsReader(flows_path, episode_uuid=uuid) as reader:
        clip_key = next(
            (k for k in reader.clip_keys() if camera_serial in reader.camera_serials(k)), None
        )
        if clip_key is None:
            raise DataError(
                f"no scene-flow clip for camera serial {camera_serial!r} in {flows_path}"
            )
        clip = reader.read_clip(clip_key, camera_serial)

    if clip.initial_rgb is not None:
        height, width = clip.initial_rgb.shape[:2]
    else:
        width, height = _FALLBACK_ANNOTATION_SIZE
        logger.warning(
            "clip %s has no initial_rgb; assuming annotation resolution %dx%d",
            clip_key, width, height,
        )
    return CameraIntrinsics.from_matrix(clip.intrinsic, width=width, height=height)


def main() -> int:
    args = parse_args()
    setup_logging()
    logger.info("episode=%s camera=%s fk_only=%s", args.episode, args.camera, args.fk_only)

    if not args.urdf.exists():
        raise DataError(
            f"URDF not found: {args.urdf}. Fetch it with "
            "`python scripts/fetch_robot_description.py --source pointworld`, or pass "
            "--urdf to point at a different robot description."
        )

    metadata = _load_metadata(args.episode)
    serial = metadata.get(f"{args.camera}_cam_serial")
    if serial is None:
        raise DataError(f"metadata for {args.episode!r} has no {args.camera}_cam_serial entry")

    mp4_path = _episode_dir(args.episode) / "recordings" / "MP4" / f"{serial}.mp4"
    if not mp4_path.exists():
        raise DataError(f"mp4 not found for camera {serial!r}: {mp4_path}")

    trajectory_path = _episode_dir(args.episode) / "trajectory.h5"
    if not trajectory_path.exists():
        raise DataError(f"trajectory.h5 not found: {trajectory_path}")

    with h5py.File(trajectory_path, "r") as traj:
        joint_positions = np.asarray(traj["observation/robot_state/joint_positions"])  # (T, 7)
        gripper_position = np.asarray(traj["observation/robot_state/gripper_position"])  # (T,)
        extrinsics_key = f"observation/camera_extrinsics/{serial}_left"
        if extrinsics_key not in traj:
            raise DataError(f"trajectory.h5 has no dataset {extrinsics_key!r}")
        extrinsics_6d = np.asarray(traj[extrinsics_key])  # (T, 6)

    n_steps = joint_positions.shape[0]
    logger.info("trajectory has %d steps", n_steps)

    base_intrinsics = _read_base_intrinsics(args.episode, serial)

    mp4_fps, _frame_count, (video_w, video_h) = read_mp4_properties(mp4_path)
    mp4_fps = mp4_fps or 15.0

    scaled_intrinsics = base_intrinsics.scaled(video_w, video_h)
    logger.info(
        "intrinsics %dx%d -> %dx%d (fx=%.1f fy=%.1f)",
        base_intrinsics.width, base_intrinsics.height, video_w, video_h,
        scaled_intrinsics.fx, scaled_intrinsics.fy,
    )

    outputs_dir = args.outputs or (REPO_ROOT / "outputs" / args.episode / "robot")
    outputs_dir.mkdir(parents=True, exist_ok=True)

    # Deferred: fpgm.robot.urdf is owned by another workstream and may not exist
    # yet at import time even though this script is written against its API.
    from fpgm.robot.urdf import RobotModel

    robot = RobotModel(str(args.urdf), load_meshes=not args.fk_only)

    n_target = n_steps if args.max_frames is None else min(n_steps, args.max_frames)

    renderer = None
    overlay_writer = None
    mask_writer = None
    packed_masks: list[np.ndarray] = []
    coverage_sum = 0.0
    n_written = 0

    if not args.fk_only:
        from fpgm.robot.render import RobotRenderer

        renderer = RobotRenderer(robot.visual_meshes(), video_w, video_h)
        overlay_writer = VideoWriter(outputs_dir / "robot_overlay.mp4", fps=mp4_fps)
        mask_writer = VideoWriter(outputs_dir / "robot_mask.mp4", fps=mp4_fps)
    else:
        overlay_writer = VideoWriter(outputs_dir / "robot_fk_dots.mp4", fps=mp4_fps)

    try:
        for t, frame_bgr in enumerate(read_frames_bgr(mp4_path)):
            if n_written >= n_target:
                break
            if t >= n_steps:
                break
            if t % args.stride != 0:
                continue

            world_to_cam = cam2world_vector_to_world2cam(extrinsics_6d[t])
            camera = Camera(scaled_intrinsics, world_to_cam)
            joints = joint_positions[t]
            gripper = float(gripper_position[t])

            if args.fk_only:
                names, origins = robot.link_origins(joints, gripper)
                uv, depth = camera.project(origins)
                visible = depth > 0
                out = draw_link_dots(
                    frame_bgr,
                    uv[visible],
                    [n for n, v in zip(names, visible, strict=True) if v],
                )
                overlay_writer.write(out)
            else:
                link_poses = robot.link_poses(joints, gripper)
                result = renderer.render(link_poses, camera)
                out = composite(frame_bgr, result)
                overlay_writer.write(out)
                mask_writer.write(mask_to_bgr(result.mask))
                packed_masks.append(np.packbits(result.mask))
                coverage = float(result.mask.mean())
                coverage_sum += coverage
                if n_written % 20 == 0:
                    logger.info("frame %d: mask coverage %.2f%%", t, coverage * 100.0)

            n_written += 1
    finally:
        overlay_writer.close()
        if mask_writer is not None:
            mask_writer.close()
        if renderer is not None:
            renderer.close()

    if not args.fk_only and packed_masks:
        npz_path = outputs_dir / "robot_mask.npz"
        np.savez_compressed(
            npz_path,
            packed=np.stack(packed_masks, axis=0),
            shape=np.array([len(packed_masks), video_h, video_w]),
        )
        logger.info("wrote %s", npz_path)

    print("\n" + "=" * 78)
    print(f"episode:        {args.episode}")
    print(f"camera:         {args.camera} ({serial})")
    print(f"frames written: {n_written}")
    if not args.fk_only and n_written:
        print(f"mean mask coverage: {100.0 * coverage_sum / n_written:.2f}%")
    print(f"outputs: {outputs_dir}")
    print("=" * 78)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())