scripts/render_robot_over_moge.py

#!/usr/bin/env python
"""VACE control = MoGe point cloud (env + objects) + URDF-rendered Franka (robot).

The split is deliberate and follows a measured failure. Feeding a MoGe-only control
(``export_moge_pointcloud_control.py``) into VACE-1.3B preserved the whole scene --
bench, towel, plate, pot, marker, all photoreal and in place -- but **replaced the
robot with a yellow toy arm** from about frame 50. Same family as this project's
earlier finding that a shaded-grey robot probe drew a *blue* toy arm: the control
carries no evidence of the robot's material, so the model invents one. MoGe cannot
supply that evidence -- it does not know a Franka is white and a Robotiq is black.
The URDF does, via its COLLADA visuals. So: geometry and appearance of everything
that is not the robot come from MoGe; the robot comes from the URDF render.

**No SAM3 anywhere.** The datagen S2 stage runs SAM3 only to *verify* the extrinsics
choice against a text-prompted mask, and on this project's own demo episode that gate
FAILED and marked every buffer UNVERIFIED anyway. This script skips it and uses
PointWorld's optimised extrinsics directly (for this episode: ``optimization_success:
true``, final loss 0.0625 on serial 20521388).

**Scale alignment, and why the robot is the right anchor.** MoGe is run without
``fov_x``, so it estimates its own FOV and its own metric scale per frame (drift CV
~9.6% on this project's demo episode) -- its depth is NOT in the same metric frame as
the URDF render's. A per-frame robust affine fit ``a*moge + b -> robot_depth`` is
solved on exactly the pixels the URDF render covers, then applied to the whole frame.
The robot is the only object that appears, metrically known, in both sources, which
makes it the one available anchor. The fit's residual is reported per frame: if it is
large, the composite's occlusion reasoning is not to be trusted and the number says so
rather than the failure being silent.

Usage:
    PYTHONPATH=src /home/quang/miniconda3/envs/fpgm/bin/python scripts/render_robot_over_moge.py \\
        --uuid RAIL+80edfcb1+2023-04-26-15h-45m-39s --camera-serial 20521388 \\
        --moge-dir outputs/moge/sample_bucket/RAIL+80edfcb1+2023-04-26-15h-45m-39s \\
        --out outputs/zeroshot_moge/RAIL+80edfcb1+2023-04-26-15h-45m-39s_urdf
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

import cv2
import h5py
import numpy as np

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

_FFMPEG = Path("/home/quang/miniconda3/envs/ffmpeg_libs/bin/ffmpeg")
_ERODE = np.ones((9, 9), np.uint8)


def write_h264(frames_bgr, out: Path, fps: float, crf: int = 14) -> None:
    it = iter(frames_bgr)
    first = next(it)
    h, w = first.shape[:2]
    out.parent.mkdir(parents=True, exist_ok=True)
    p = subprocess.Popen(
        [str(_FFMPEG), "-y", "-hide_banner", "-loglevel", "error", "-f", "rawvideo",
         "-pix_fmt", "bgr24", "-s", f"{w}x{h}", "-r", f"{fps:.4f}", "-i", "pipe:0",
         "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", str(crf), str(out)],
        stdin=subprocess.PIPE)
    assert p.stdin is not None
    try:
        p.stdin.write(np.ascontiguousarray(first).tobytes())
        for f in it:
            p.stdin.write(np.ascontiguousarray(f).tobytes())
    finally:
        p.stdin.close()
        if p.wait() != 0:
            raise RuntimeError(f"ffmpeg failed writing {out}")


def robust_scale(src: np.ndarray, dst: np.ndarray) -> tuple[float, float]:
    """Median of ``dst/src`` -- a SCALE-ONLY fit. Returns ``(a, residual_m)``.

    A scale+shift (affine) fit was tried first and is wrong here, measured: on the
    robot mask alone the depth range is narrow, so ``a`` and ``b`` trade off freely
    and ``a`` wandered between 0.037 and 1.08 across frames (CV 0.84) -- the fit was
    ill-conditioned, not the extrinsics. Dropping the shift makes it stable at
    CV 0.017. That is also the physically right model: MoGe-2 predicts *metric*
    depth (unlike affine-invariant MiDaS-family models), so the only free parameter
    its per-frame FOV guess can introduce is a scale.
    """
    r = dst / np.maximum(src, 1e-6)
    a = float(np.median(r))
    return a, float(np.median(np.abs(a * src - dst)))


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--uuid", required=True)
    ap.add_argument("--camera-serial", required=True)
    ap.add_argument("--moge-dir", type=Path, required=True)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--config", type=Path, default=REPO_ROOT / "configs/datagen_droid.yaml")
    ap.add_argument("--fps", type=float, default=15.0)
    ap.add_argument("--depth-tol-m", type=float, default=0.02,
                    help="robot wins where robot_depth < aligned_moge + tol")
    args = ap.parse_args()

    from fpgm.config_datagen import DatagenProfile
    from fpgm.datagen.frame_index import EpisodeFrameIndex
    from fpgm.data.pointworld import FlowsReader
    from fpgm.datagen.robot_buffers import (
        load_extrinsics_candidates, read_native_camera_intrinsics,
    )
    from fpgm.geometry.camera import Camera
    from fpgm.robot.render import RobotRenderer
    from fpgm.robot.urdf import RobotModel
    from fpgm.viz.overlays import read_frames_bgr

    profile = DatagenProfile.from_yaml(str(args.config))
    traj = profile.paths.droid_episode_dir(args.uuid) / "trajectory.h5" \
        if hasattr(profile.paths, "droid_episode_dir") else \
        REPO_ROOT / "data/droid_raw" / args.uuid / "trajectory.h5"
    mp4 = next((REPO_ROOT / "data/droid_raw" / args.uuid / "recordings/MP4").glob(
        f"{args.camera_serial}.mp4"))

    cap = cv2.VideoCapture(str(mp4))
    video_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    video_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    n_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    mp4_fps = float(cap.get(cv2.CAP_PROP_FPS))
    cap.release()

    flows = profile.paths.flows_h5(args.uuid)
    with FlowsReader(flows, episode_uuid=args.uuid) as reader:
        frame_index = EpisodeFrameIndex.build(
            args.uuid, reader, traj,
            mp4_properties=(mp4_fps, n_video_frames, (video_w, video_h)),
            camera_serial=args.camera_serial)

    intr = read_native_camera_intrinsics(flows, args.camera_serial).scaled(video_w, video_h)
    cands = load_extrinsics_candidates(
        profile.paths.cameras_json(args.uuid), traj, args.camera_serial)
    chosen = next(c for c in cands if c.name == "optimized_cameras_json")
    camera = Camera(intr, chosen.world_to_cam)
    print(f"extrinsics: {chosen.name} | intrinsics {intr.width}x{intr.height} "
          f"fx={intr.fx:.1f} fy={intr.fy:.1f}")

    robot = RobotModel(str(profile.paths.urdf), load_meshes=True)
    with h5py.File(traj, "r") as f:
        from fpgm.datagen.robot_buffers import (
            TRAJECTORY_GRIPPER_POSITION_KEY, TRAJECTORY_JOINT_POSITIONS_KEY,
        )
        joint_positions = np.asarray(f[TRAJECTORY_JOINT_POSITIONS_KEY])
        gripper = np.asarray(f[TRAJECTORY_GRIPPER_POSITION_KEY])
    renderer = RobotRenderer(robot.visual_meshes(), intr.width, intr.height)

    npz_files = sorted((args.moge_dir / "per_frame").glob("frame_*.npz"))
    n = min(n_video_frames, len(npz_files))
    print(f"{args.uuid}: {n} frames, {video_w}x{video_h}")

    args.out.mkdir(parents=True, exist_ok=True)
    stats = {"scale_a": [], "resid_m": [], "robot_px": []}
    last_a = [1.0]

    def frames():
        for t, frame_bgr in enumerate(read_frames_bgr(mp4)):
            if t >= n:
                break
            row = frame_index.trajectory_row(t)
            rr = renderer.render(
                robot.link_poses(joint_positions[row], float(gripper[row])), camera)

            d = np.load(npz_files[t])
            moge = d["depth"].astype(np.float32)
            moge = np.where(d["mask"].astype(bool), moge, np.nan)

            # Erode the mask before fitting: silhouette pixels have the render's
            # depth on the robot but MoGe's on whatever is behind it, and they are
            # pure outliers to a depth ratio.
            core = cv2.erode(rr.mask.astype(np.uint8), _ERODE).astype(bool)
            fit_px = core & np.isfinite(moge) & (rr.depth > 0)
            if fit_px.sum() >= 200:
                a, resid = robust_scale(moge[fit_px], rr.depth[fit_px])
                last_a[0] = a
            else:
                a, resid = last_a[0], float("nan")  # hold the last good scale
            stats["scale_a"].append(a)
            stats["resid_m"].append(resid)
            stats["robot_px"].append(int(rr.mask.sum()))

            moge_aligned = a * moge
            # Robot wins only where it is actually in front of the aligned MoGe
            # surface -- so an arm passing behind a foreground object stays occluded
            # instead of being pasted over it.
            win = rr.mask & (rr.depth > 0) & (
                ~np.isfinite(moge_aligned) | (rr.depth < moge_aligned + args.depth_tol_m))
            out = frame_bgr.copy()
            out[win] = rr.color[win][:, ::-1]  # RenderResult.color is RGB
            yield out

    out_video = args.out / "control_moge_urdf.mp4"
    write_h264(frames(), out_video, args.fps)
    renderer.close()

    ref = args.out / "ref_frame0.png"
    cap = cv2.VideoCapture(str(mp4))
    ok, f0 = cap.read()
    cap.release()
    cv2.imwrite(str(ref), f0)

    a = np.array(stats["scale_a"]); r = np.array(stats["resid_m"])
    meta = {"uuid": args.uuid, "camera_serial": args.camera_serial,
            "extrinsics": chosen.name, "n_frames": n,
            "scale_a_median": round(float(np.median(a)), 4),
            "scale_a_p10_p90": [round(float(np.percentile(a, 10)), 4),
                                round(float(np.percentile(a, 90)), 4)],
            "align_residual_m_median": round(float(np.nanmedian(r)), 4),
            "robot_px_median": int(np.median(stats["robot_px"])),
            "control": str(out_video), "ref_image": str(ref)}
    (args.out / "control_meta.json").write_text(json.dumps(meta, indent=2))
    print(json.dumps(meta, indent=2))


if __name__ == "__main__":
    main()