scripts/render_fixed_initial_view.py

#!/usr/bin/env python
"""Render robot+object motion held to the INITIAL frame's camera + background.

Unlike `run_object_pipeline.py`'s `render` stage -- which either recomposites
onto each frame of the *real* video (`scene_overlay.mp4`, camera moves because
the real camera moved) or orbits a synthetic freecam (`scene_freecam.mp4`) --
this holds both the camera pose and the background image fixed at frame 0 of
the clip. Only the robot (FK per frame) and the object (per-frame pose from
`ObjectTrajectory`) move; the frame-0 point cloud is naturally static already.
This is the "initial frame + moving robot + moving object" composite the
freecam/overlay renders don't give directly.

Reuses `scripts/run_object_pipeline.py`'s already-computed per-clip artifacts
(mesh/align/act stages must already have been run for --outputs) rather than
recomputing anything -- this is a render-only pass.

Usage:
    PYTHONPATH=src python scripts/render_fixed_initial_view.py \\
        --episode <uuid> --camera ext1 --clip 20:31 \\
        --outputs outputs/<uuid>/objects_sam3d_multi/20_31 \\
        --urdf third_party/robot_description/pointworld_franka_robotiq_2f85/franka_panda_robotiq_2f85.urdf \\
        --out outputs/fixed_view_20_31.mp4
"""

from __future__ import annotations

import argparse
import importlib.util
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.geometry.camera import Camera  # noqa: E402
from fpgm.objects.scene import SceneRenderer, draw_point_cloud, draw_scene_hud, project_point_cloud  # noqa: E402
from fpgm.pipeline.frames import ClipFrameSource  # noqa: E402
from fpgm.robot.overlay import composite  # noqa: E402
from fpgm.robot.urdf import RobotModel  # noqa: E402
from fpgm.viz.overlays import VideoWriter  # noqa: E402

# run_object_pipeline.py is a script, not a package -- import it by path so its
# private `_load_mesh`/`_load_alignment`/`_load_trajectory`/`_load_clip_context`
# helpers (the exact readers that already know each stage's on-disk layout)
# can be reused instead of re-implementing them here.
_spec = importlib.util.spec_from_file_location(
    "run_object_pipeline", REPO_ROOT / "scripts" / "run_object_pipeline.py"
)
rop = importlib.util.module_from_spec(_spec)
sys.modules[_spec.name] = rop  # dataclass() introspects sys.modules[cls.__module__] at class-body time
_spec.loader.exec_module(rop)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--episode", required=True)
    p.add_argument("--camera", default="ext1", choices=["ext1", "ext2"])
    p.add_argument("--clip", required=True, help='PointWorld clip key, e.g. "20:31"')
    p.add_argument("--outputs", type=Path, required=True, help="stage outputs root (mesh/align/act already run)")
    p.add_argument("--urdf", type=Path, required=True)
    p.add_argument("--out", type=Path, required=True)
    return p.parse_args()


def main() -> int:
    args = parse_args()

    mesh_dir = rop.stage_dir(args.outputs, "mesh")
    align_dir = rop.stage_dir(args.outputs, "align")
    act_dir = rop.stage_dir(args.outputs, "act")

    meta = rop._load_mask_meta(args.outputs)
    mesh = rop._load_mesh(mesh_dir)
    alignment = rop._load_alignment(align_dir)
    trajectory = rop._load_trajectory(act_dir)
    n = len(trajectory)
    if n == 0:
        raise SystemExit("empty trajectory -- nothing to render")

    robot = RobotModel(str(args.urdf), load_meshes=True)

    ctx = rop._load_clip_context(args.episode, args.camera, args.clip)
    video_w, video_h = meta["video_resolution"]
    camera_video = ctx.camera_annot.rescaled(video_w, video_h)

    traj_path = rop._load_trajectory_h5(args.episode)
    with h5py.File(traj_path, "r") as f:
        joint_positions_full = np.asarray(f["observation/robot_state/joint_positions"])
        gripper_full = np.asarray(f["observation/robot_state/gripper_position"])
        extrinsics_key = f"observation/camera_extrinsics/{meta['camera_serial']}_left"
        extrinsics_6d_full = np.asarray(f[extrinsics_key])

    stride = ctx.timing.annotation_stride
    joint_rows = rop._strided_rows(ctx.clip.start, n, stride, joint_positions_full.shape[0], "render (joints)")
    extrinsics_rows = rop._strided_rows(ctx.clip.start, n, stride, extrinsics_6d_full.shape[0], "render (extrinsics)")
    joint_positions = joint_positions_full[joint_rows]
    gripper = gripper_full[joint_rows]
    extrinsics_6d = extrinsics_6d_full[extrinsics_rows]

    mp4_fps, _count, _size = read_mp4_properties(ctx.mp4_path)
    fps = mp4_fps or rop._TRAJECTORY_FPS

    point_colors = (
        alignment.point_colors
        if alignment.point_colors is not None
        else np.zeros_like(alignment.points_world, dtype=np.uint8)
    )

    # The whole point: camera + background frozen at frame 0 -- everything
    # below reuses index 0 for both, while robot/object poses still index `t`.
    fixed_world_to_cam = cam2world_vector_to_world2cam(extrinsics_6d[0])
    fixed_camera = Camera(camera_video.K, fixed_world_to_cam)

    video_start = int(ctx.timing.clip_frame_to_video_frame(0))
    video_end = int(ctx.timing.clip_frame_to_video_frame(n))

    scene_renderer = SceneRenderer(robot.visual_meshes(), mesh, video_w, video_h)
    try:
        with ClipFrameSource(str(ctx.mp4_path), video_start, video_end, stride=stride) as frames, \
                VideoWriter(args.out, fps=fps) as writer:
            frame0_bgr = frames.read(0).copy()
            speeds = rop._object_speeds(trajectory)
            for t in range(n):
                link_poses = robot.link_poses(joint_positions[t], float(gripper[t]))
                object_pose = trajectory.transforms[t]
                state = trajectory.states[t]
                timestamp = float(trajectory.timestamps[t])

                result = scene_renderer.render(link_poses, object_pose, fixed_camera)
                composed = composite(frame0_bgr, result, alpha=0.85)
                uv, cols = project_point_cloud(alignment.points_world, point_colors, fixed_camera, (video_h, video_w))
                composed = draw_point_cloud(composed, uv, cols, radius=2)
                composed = draw_scene_hud(composed, t, timestamp, state, float(speeds[t]), trajectory.push_gain)
                writer.write(composed)
    finally:
        scene_renderer.close()

    print(str(args.out.resolve()))
    return 0


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