scripts/render_pointcloud_scene.py

#!/usr/bin/env python
"""Render the scene as a depth-derived point cloud, with robot/objects as meshes.

Built for use as a 3D-geometry conditioning signal for video diffusion: the
static scene is carried by the *measured* point cloud (so the conditioning
encodes real metric geometry rather than a photograph), while the robot and the
manipulated objects -- the parts that move, and whose motion is known exactly
from forward kinematics and the object trajectory -- are rendered as solid
meshes on top.

The point cloud comes from PointWorld's dense `initial_depth` for the clip
(320x180, ~91% valid here), unprojected through that clip's intrinsics into the
world/robot-base frame and coloured from the paired `initial_rgb`. It is a
single-viewpoint capture, so it is a 2.5-D shell: surfaces the camera could not
see have no points, and orbiting far from the capture viewpoint exposes those
holes. That is a property of the data, not a bug to paper over -- for
conditioning at or near the recorded viewpoint it is exactly right.

Occlusion between the cloud and the meshes is resolved against the renderer's
own depth buffer, not by draw order: a point is drawn only where it is nearer
than whatever mesh surface covers that pixel, so the arm correctly hides the
scene behind it.

Usage:
    PYTHONPATH=src python scripts/render_pointcloud_scene.py \\
        --episode <uuid> --drawer-outputs <...> --brick-outputs <...> \\
        --calibration outputs/push_calibration_tracked.json \\
        --mode fixed --out outputs/pointcloud_scene.mp4
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import sys
from pathlib import Path

import h5py
import numpy as np
import trimesh
from PIL import Image

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  # noqa: E402
from fpgm.geometry.camera import Camera  # noqa: E402
from fpgm.objects.scene import free_camera  # noqa: E402
from fpgm.robot.render import RobotRenderer  # noqa: E402
from fpgm.robot.urdf import RobotModel  # noqa: E402
from fpgm.types import CameraIntrinsics  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402
from fpgm.viz.overlays import VideoWriter, draw_hud  # noqa: E402

import pyrender  # noqa: E402

_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
_spec.loader.exec_module(rop)

logger = get_logger("pointcloud_scene")

_TRAJ_FPS = 15.0
_BG = (12, 12, 16)


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", default="0:11", help="clip whose initial_depth/initial_rgb build the cloud")
    p.add_argument("--drawer-outputs", type=Path, default=None)
    p.add_argument("--brick-outputs", type=Path, default=None)
    p.add_argument("--calibration", type=Path, default=None,
                   help="push measurement; drives the drawer's per-frame slide")
    p.add_argument("--contact-rows", type=int, nargs=2, default=[89, 108])
    p.add_argument("--start-row", type=int, default=81)
    p.add_argument("--mode", choices=["fixed", "orbit"], default="fixed")
    p.add_argument("--orbit-degrees", type=float, default=50.0)
    p.add_argument("--width", type=int, default=1280)
    p.add_argument("--height", type=int, default=720)
    p.add_argument("--point-radius", type=int, default=2)
    p.add_argument("--max-depth", type=float, default=3.0, help="drop cloud points beyond this (metres)")
    p.add_argument("--n-frames", type=int, default=60)
    p.add_argument("--capture-row", type=int, default=0,
                   help="trajectory row the depth frame corresponds to (video frame index)")
    p.add_argument("--keep-dynamic-points", action="store_true",
                   help="do NOT carve the robot/objects out of the cloud (they would appear twice)")
    p.add_argument("--robot-dilate-px", type=int, default=2,
                   help="dilate the robot silhouette before carving, to catch depth bleed at edges")
    p.add_argument("--object-pad-m", type=float, default=0.02,
                   help="padding around each object's bbox when carving it out of the cloud")
    p.add_argument("--urdf", type=Path, default=None)
    p.add_argument("--out", type=Path, required=True)
    return p.parse_args()


class _MultiObjectRenderer:
    """RobotRenderer plus N independently posed object meshes (see
    scripts/render_counterfactual_push.py for the same pattern)."""

    def __init__(self, link_meshes, object_meshes, width, height, **kw):
        self._robot = RobotRenderer(link_meshes, width, height, **kw)
        self._nodes = []
        for tri in object_meshes:
            node = pyrender.Node(mesh=pyrender.Mesh.from_trimesh(tri, smooth=False), matrix=np.eye(4))
            self._robot.scene.add_node(node)
            self._nodes.append(node)

    def render(self, link_poses, object_poses, camera):
        for node, pose in zip(self._nodes, object_poses, strict=True):
            self._robot.scene.set_pose(node, np.asarray(pose, dtype=np.float64))
        return self._robot.render(link_poses, camera)

    def close(self):
        self._robot.close()


def _scene_point_cloud(ctx, depth_u16, max_depth):
    """Unproject every valid depth pixel into world-frame points + RGB colours."""
    depth_m = depth_u16.astype(np.float64) / 1000.0
    h, w = depth_m.shape
    vs, us = np.nonzero((depth_m > 0) & (depth_m <= max_depth))
    z = depth_m[vs, us]
    uv = np.stack([us + 0.5, vs + 0.5], axis=1).astype(np.float64)
    pts = ctx.camera_annot.unproject(uv, z)

    rgb = ctx.clip.initial_rgb
    if rgb is not None and rgb.shape[:2] == (h, w):
        cols = rgb[vs, us]
    else:
        cols = np.full((pts.shape[0], 3), 200, dtype=np.uint8)
    return pts, cols.astype(np.uint8)


def _draw_points(canvas, uv, depth, colors, mesh_depth, radius):
    """Paint cloud points, nearest last, hidden where a mesh surface is nearer.

    Vectorised on purpose: a per-point cv2.circle loop over ~50k points x tens
    of frames dominates runtime, whereas fancy-indexed assignment in
    far-to-near order gives the same painter's-algorithm result in one pass per
    stamp offset.
    """
    h, w = canvas.shape[:2]
    u = np.round(uv[:, 0]).astype(np.int64)
    v = np.round(uv[:, 1]).astype(np.int64)
    keep = (depth > 0) & (u >= 0) & (u < w) & (v >= 0) & (v < h)
    u, v, depth, colors = u[keep], v[keep], depth[keep], colors[keep]

    # Hide points behind rendered geometry (mesh_depth == 0 means "no mesh").
    md = mesh_depth[v, u]
    visible = (md <= 0.0) | (depth < md - 1e-3)
    u, v, depth, colors = u[visible], v[visible], depth[visible], colors[visible]

    order = np.argsort(depth)[::-1]  # far first so nearer points overwrite
    u, v, colors = u[order], v[order], colors[order]
    bgr = colors[:, ::-1]

    r = int(radius)
    for du in range(-r, r + 1):
        for dv in range(-r, r + 1):
            if du * du + dv * dv > r * r:
                continue
            uu = np.clip(u + du, 0, w - 1)
            vv = np.clip(v + dv, 0, h - 1)
            canvas[vv, uu] = bgr
    return canvas


def _drop_robot_points(cloud_xyz, cloud_rgb, robot, joints, gripper, ctx, world_to_cam, dilate_px):
    """Remove cloud points that are the ROBOT as the depth camera saw it.

    The cloud is a snapshot of the capture frame, so it contains the arm frozen
    in that one pose. Left in, it renders as a static ghost arm while the mesh
    arm moves away from it -- the same surface represented twice, which is
    exactly the wrong signal for geometry conditioning. Removal is done by
    rendering the robot at the capture pose through the capture camera and
    dropping every cloud point whose own reprojection lands on that silhouette,
    so it keys on what the camera actually saw rather than on a hand-tuned
    bounding volume.
    """
    import cv2

    w, h = ctx.camera_annot.K.width, ctx.camera_annot.K.height
    cam = Camera(ctx.camera_annot.K, world_to_cam)
    rr = RobotRenderer(robot.visual_meshes(), w, h)
    try:
        res = rr.render(robot.link_poses(joints, gripper), cam)
        mask = res.mask.copy()
    finally:
        rr.close()

    if dilate_px > 0:
        k = np.ones((2 * dilate_px + 1, 2 * dilate_px + 1), np.uint8)
        mask = cv2.dilate(mask.astype(np.uint8), k).astype(bool)

    uv, dep = cam.project(cloud_xyz)
    u = np.clip(np.round(uv[:, 0]).astype(int), 0, w - 1)
    v = np.clip(np.round(uv[:, 1]).astype(int), 0, h - 1)
    on_robot = mask[v, u] & (dep > 0)
    logger.info("robot removal: dropped %d/%d cloud points", int(on_robot.sum()), cloud_xyz.shape[0])
    keep = ~on_robot
    return cloud_xyz[keep], cloud_rgb[keep]


def _drop_object_points(cloud_xyz, cloud_rgb, tri, transform, pad, label):
    """Remove cloud points lying inside a posed object's own bounding box.

    Done in the object's canonical frame (points pulled back through its
    alignment) rather than in world axes, so a rotated object does not need an
    inflated axis-aligned box that would eat the surrounding scene.
    """
    R = transform[:3, :3]
    t = transform[:3, 3]
    scale = np.linalg.norm(R, axis=0)
    scale[scale < 1e-9] = 1.0
    rot = R / scale[None, :]
    local = (cloud_xyz - t) @ rot
    local = local / scale[None, :]

    lo = np.asarray(tri.vertices).min(axis=0) - pad / scale
    hi = np.asarray(tri.vertices).max(axis=0) + pad / scale
    inside = np.all((local >= lo) & (local <= hi), axis=1)
    logger.info("%s removal: dropped %d/%d cloud points", label, int(inside.sum()), cloud_xyz.shape[0])
    keep = ~inside
    return cloud_xyz[keep], cloud_rgb[keep]


def _load_object(stage_root: Path):
    mesh = rop._load_mesh(rop.stage_dir(stage_root, "mesh"))
    alignment = rop._load_alignment(rop.stage_dir(stage_root, "align"))
    # process=False: welding vertices here would desync a textured mesh's
    # per-vertex `uv` from the vertex it belongs to (see
    # scripts/run_object_pipeline.py's _save_mesh/_load_mesh docstrings).
    if mesh.uv is not None and mesh.texture is not None:
        visual = trimesh.visual.TextureVisuals(
            uv=np.asarray(mesh.uv, dtype=np.float32),
            image=Image.fromarray(np.asarray(mesh.texture, dtype=np.uint8)),
        )
        tri = trimesh.Trimesh(
            vertices=np.asarray(mesh.vertices, dtype=np.float64),
            faces=np.asarray(mesh.faces, dtype=np.int64),
            visual=visual,
            process=False,
        )
        # No vertex_colors assignment here: it would replace `visual` with a
        # fresh ColorVisuals and silently drop the texture (trimesh visuals
        # are one-kind-at-a-time) -- the exact bug this whole fix removes.
    else:
        tri = trimesh.Trimesh(
            vertices=np.asarray(mesh.vertices, dtype=np.float64),
            faces=np.asarray(mesh.faces, dtype=np.int64),
            process=False,
        )
        if mesh.vertex_colors is not None:
            tri.visual.vertex_colors = mesh.vertex_colors
    return tri, alignment.transform


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

    ctx = rop._load_clip_context(args.episode, args.camera, args.clip)
    metadata = rop._load_metadata(args.episode)
    serial = metadata[f"{args.camera}_cam_serial"]
    depth_u16 = rop._load_initial_depth(args.episode, args.clip, serial)
    cloud_xyz, cloud_rgb = _scene_point_cloud(ctx, depth_u16, args.max_depth)
    logger.info("scene cloud: %d points", cloud_xyz.shape[0])

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

    object_meshes, object_T0 = [], []
    for root in (args.drawer_outputs, args.brick_outputs):
        if root is None:
            continue
        tri, T = _load_object(root)
        object_meshes.append(tri)
        object_T0.append(T)

    with h5py.File(rop._load_trajectory_h5(args.episode), "r") as f:
        J = np.asarray(f["observation/robot_state/joint_positions"])
        G = np.asarray(f["observation/robot_state/gripper_position"])
        extrinsics_6d = np.asarray(f[f"observation/camera_extrinsics/{serial}_left"])

    # Everything that will be drawn as a mesh must first be carved out of the
    # cloud, or it appears twice: once frozen at the capture pose, once moving.
    if not args.keep_dynamic_points:
        capture_w2c = cam2world_vector_to_world2cam(extrinsics_6d[args.capture_row])
        cloud_xyz, cloud_rgb = _drop_robot_points(
            cloud_xyz, cloud_rgb, robot, J[args.capture_row], float(G[args.capture_row]),
            ctx, capture_w2c, args.robot_dilate_px,
        )
        for tri, T, label in zip(object_meshes, object_T0,
                                 ["drawer", "brick"][: len(object_meshes)], strict=False):
            cloud_xyz, cloud_rgb = _drop_object_points(
                cloud_xyz, cloud_rgb, tri, T, args.object_pad_m, label
            )
        logger.info("scene cloud after removal: %d points", cloud_xyz.shape[0])

    # Drawer slide, replayed from the measured push when a calibration is given.
    r0, r1 = args.contact_rows
    slide_axis = np.zeros(3)
    travel = None
    if args.calibration is not None and args.drawer_outputs is not None:
        cal = json.loads(args.calibration.read_text())
        disp = np.array(
            [[np.nan, np.nan] if v is None or v[0] is None else v for v in cal["displacement_xy_m"]],
            dtype=float,
        )
        axis2 = disp[r1] / (np.linalg.norm(disp[r1]) + 1e-12)
        slide_axis = np.array([axis2[0], axis2[1], 0.0])
        travel = np.linalg.norm(np.nan_to_num(disp), axis=1)

    K = CameraIntrinsics(
        fx=ctx.camera_annot.K.fx * args.width / ctx.camera_annot.K.width,
        fy=ctx.camera_annot.K.fy * args.height / ctx.camera_annot.K.height,
        cx=ctx.camera_annot.K.cx * args.width / ctx.camera_annot.K.width,
        cy=ctx.camera_annot.K.cy * args.height / ctx.camera_annot.K.height,
        width=args.width, height=args.height,
    )
    base_cam = Camera(K, cam2world_vector_to_world2cam(extrinsics_6d[args.start_row]))
    cloud_centre = cloud_xyz.mean(axis=0)
    orbit_distance = float(np.linalg.norm(base_cam.cam_to_world(np.zeros((1, 3)))[0] - cloud_centre))

    renderer = _MultiObjectRenderer(robot.visual_meshes(), object_meshes, args.width, args.height,
                                   bg_color=(0.0, 0.0, 0.0, 0.0))
    rows = np.linspace(args.start_row, r1, args.n_frames)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    try:
        with VideoWriter(args.out, fps=30.0) as writer:
            for i, row_f in enumerate(rows):
                row = int(round(row_f))
                link_poses = robot.link_poses(J[row], float(G[row]))

                poses = []
                for j, T in enumerate(object_T0):
                    pose = T.copy()
                    if j == 0 and travel is not None:
                        pose[:3, 3] = pose[:3, 3] + float(travel[row] - travel[r0]) * slide_axis \
                            if row >= r0 else pose[:3, 3]
                    poses.append(pose)

                if args.mode == "orbit":
                    az = -args.orbit_degrees / 2 + args.orbit_degrees * i / max(args.n_frames - 1, 1)
                    camera = free_camera(cloud_centre, orbit_distance, 180.0 + az, 18.0, K)
                else:
                    camera = base_cam

                result = renderer.render(link_poses, poses, camera)
                canvas = np.full((args.height, args.width, 3), _BG, dtype=np.uint8)
                uv, dep = camera.project(cloud_xyz)
                canvas = _draw_points(canvas, uv, dep, cloud_rgb, result.depth, args.point_radius)

                mask = result.mask
                canvas[mask] = result.color[:, :, ::-1][mask]
                canvas = draw_hud(canvas, [
                    f"scene = {cloud_xyz.shape[0]} depth points   robot+objects = meshes",
                    f"frame {i+1}/{args.n_frames}   traj row {row}   mode={args.mode}",
                    "single-viewpoint 2.5D capture: unseen surfaces have no points",
                ])
                writer.write(canvas)
    finally:
        renderer.close()

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


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