scripts/appearance_control.py

#!/usr/bin/env python
"""Stage 2: (control, target) appearance-training pairs for the VACE fine-tune.

Same output shape as the zero-shot probe (``render_robot_over_moge.py``): control
= the real frame with a URDF-rendered Franka composited over the real arm, target
= the untouched real frame. That pair is exactly the appearance mapping the
fine-tune has to learn -- "this white/black CAD render is that real robot" -- and
it is the mapping VACE-1.3B provably does not have zero-shot: given a control
whose robot was already a clean, correctly-posed white Franka + black Robotiq, it
still drew a yellow toy arm from about frame 50.

**No object meshes, no SAM3, no quality gates.** The datagen pipeline's S2/S3/S6/S7
gates exist to certify object *pose and geometry*; appearance training needs
neither, so nothing here can reject an episode and there is no pass rate to lose.

**No PointWorld.** Measured, not assumed:
  * the frame->trajectory-row map is ``n_video_frames / n_trajectory_rows``
    (``ClipTiming.from_counts``); the flows reader was only ever used to enumerate
    clips, which this script does not need;
  * extrinsics come from ``cameras.json`` (172 MB for 42935 episodes, already on
    disk) with ``metadata.json`` as fallback;
  * intrinsics are a property of the physical ZED unit, not the episode -- across
    62 episodes of serial 22008760 the focal length varies by 0.011%, and all of
    DROID uses only 33 distinct exterior cameras -- so a serial->intrinsic table
    harvested once covers every episode those cameras ever recorded.
That drops the 2.28 TB flows dependency to ~0.

**Everything runs at the training resolution, and both videos are written there.**
Measured at 832x480 against 1280x720: render 56.3 vs 43.4 fps, composite 159 vs
46 fps, encode 141 vs 72 fps. Writing the target out too (rather than symlinking
the source mp4 and cropping at train time) costs one extra H.264 generation on
top of DROID's own, and buys: a self-contained training set at 2.8 MB/clip
instead of 17.9, a dataloader that does zero resizing, and no runtime dependency
on ``data/droid_raw``. The trade is that changing training resolution means
re-running this stage -- deliberate, since VACE-1.3B trains at one resolution.

**Occlusion tolerance.** The robot wins where its render is in front of the
MoGe surface by more than ``--depth-tol-m``. The default 0.15 m is measured, not
picked: on a clip with no true occlusion at all, a 0.02 m margin still rejects
18.8% of robot pixels, 0.05 m rejects 13.4%, 0.10 m rejects 4.0% and 0.20 m
rejects 0.078% -- a tight margin is not detecting occlusion, it is punching noise
holes in the arm. Across the first 126-clip run 0.15 m rejected 0.20% (median).

Env: ``fpgm``, EGL (56.3 fps at 832x480, against 2.7 fps on the OSMesa default).
Usage:
    PYOPENGL_PLATFORM=egl PYTHONPATH=src \\
    /home/quang/miniconda3/envs/fpgm/bin/python -u scripts/appearance_control.py \\
        --work <scratch> --out outputs/appearance_pairs --follow
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
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"))
sys.path.insert(0, str(REPO_ROOT / "scripts"))
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")

from appearance_depth import cover_crop, cover_crop_params  # noqa: E402

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


def ffmpeg() -> str:
    return str(_FFMPEG) if _FFMPEG.exists() else (shutil.which("ffmpeg") or "ffmpeg")


class H264Writer:
    """Streaming H.264 writer. Frames are pushed in; ffmpeg runs concurrently.

    ``crf 18`` at 832x480 measured 9.4 KB/frame against 18 KB at crf 14, for
    output that is still above the quality of the DROID source it came from
    (already H.264-compressed, then downscaled here, which removes the
    high-frequency content a low crf would have been spending bits on).
    Encoding at 141 fps it is not the bottleneck, so ``preset medium`` is kept
    rather than trading quality for speed that nothing would use.
    """

    def __init__(self, out: Path, w: int, h: int, fps: float, crf: int) -> None:
        out.parent.mkdir(parents=True, exist_ok=True)
        self.proc = subprocess.Popen(
            [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", "-preset", "medium", "-pix_fmt", "yuv420p",
             "-crf", str(crf), str(out)],
            stdin=subprocess.PIPE)
        self.out = out

    def write(self, frame_bgr: np.ndarray) -> None:
        self.proc.stdin.write(np.ascontiguousarray(frame_bgr).tobytes())

    def close(self) -> None:
        self.proc.stdin.close()
        if self.proc.wait() != 0:
            raise RuntimeError(f"ffmpeg failed writing {self.out}")


def rel(p: Path) -> str:
    """Repo-relative string, whichever way the caller spelled the path."""
    p = Path(p).resolve()
    try:
        return str(p.relative_to(REPO_ROOT))
    except ValueError:
        return str(p)


def load_intrinsic_table(path: Path) -> dict:
    return json.loads(path.read_text())


def intrinsics_for(table: dict, serial: str, w: int, h: int):
    """Real calibration for ``serial`` at the *source* ``w x h``. ``(intr, source)``.

    Falls back to the table's ``_default`` (the median ZED unit) for a serial that
    was never harvested. The fallback is recorded in the clip's meta so an audit
    can find every clip built on a guessed principal point -- across the 12 known
    units cy/h spans 0.478-0.518, which is 29 px at 720, so the guess is not free.
    """
    from fpgm.types import CameraIntrinsics

    entry = table["cameras"].get(serial)
    source = f"table:{serial}"
    if entry is None:
        entry, source = table["_default"], "default"
    return CameraIntrinsics(
        fx=entry["fx_over_w"] * w, fy=entry["fy_over_h"] * h,
        cx=entry["cx_over_w"] * w, cy=entry["cy_over_h"] * h,
        width=w, height=h,
    ), source


def crop_intrinsics(intr, out_w: int, out_h: int):
    """Rescale + crop-shift ``intr`` so it stays valid for :func:`cover_crop`'s output.

    Mirrors ``_CoverCrop.apply_intrinsics``. Without this the render would be
    projected with the uncropped camera and land in the wrong place -- the crop
    is not a centred one (832/1280 != 480/720), so it shifts cx by ~10 px.
    """
    from fpgm.types import CameraIntrinsics

    rw, rh, ox, oy = cover_crop_params(intr.width, intr.height, out_w, out_h)
    sx, sy = rw / intr.width, rh / intr.height
    return CameraIntrinsics(fx=intr.fx * sx, fy=intr.fy * sy,
                            cx=intr.cx * sx - ox, cy=intr.cy * sy - oy,
                            width=out_w, height=out_h)


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

    MoGe-2 is metric, so its per-frame FOV guess can only introduce a scale; the
    scale+shift version was measured to be ill-conditioned on the narrow depth
    range a robot mask spans (``a`` wandered 0.037-1.08, CV 0.84, against CV
    0.017 for scale-only).
    """
    r = dst / np.maximum(src, 1e-6)
    a = float(np.median(r))
    return a, float(np.median(np.abs(a * src - dst)))


def render_and_composite(
    rend, robot, joint_row: np.ndarray, gripper_value: float,
    frame_bgr: np.ndarray, moge: np.ndarray, camera, depth_tol_m: float, last_a: float,
    *, background_bgr: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, float, float]:
    """One frame of the appearance-pair composite: render, depth-align, occlude, draw.

    Extracted from :meth:`ClipBuilder.render_loop` (formerly inlined in
    ``ClipBuilder.build``) so a second caller can drive it with a gripper value
    that did not come from ``trajectory.h5`` -- see
    ``scripts/render_counterfactual_grip.py``, which renders the *same* clip
    with the recorded arm pose but an overridden gripper trace. Nothing here
    changed: this is the original loop body with its four locals promoted to
    parameters and its five locals promoted to a return tuple.

    ``background_bgr``, if given, replaces ``frame_bgr`` as the base the robot
    is painted onto -- ``frame_bgr`` itself is still required (and still drives
    the MoGe depth-alignment fit's caller-side inputs, unchanged) but is then
    otherwise unused for that frame. This exists for
    ``scripts/render_background_probe.py``: every appearance pair ever built
    (800 clips, the whole training set) uses the real camera frame as its
    background, which hands the model the table/drawer/object for free and
    leaves open whether the fine-tune can generate a scene at all, as opposed
    to only knowing how to repaint a robot onto one it was already given. A
    control with a synthetic background (a temporal-median plate, flat grey,
    flat black) answers that, without touching the 800 pairs already built on
    the real-frame path. **Default (``None``) reproduces the exact prior
    behaviour** -- ``out = frame_bgr.copy()`` -- byte for byte: this is a
    strict superset of the old signature, not a new codepath, and every
    existing caller (``ClipBuilder.render_loop`` and
    ``render_counterfactual_grip.py``) calls it positionally without this
    keyword and is therefore unaffected.

    Returns:
        ``(composited_bgr, drawn_mask, occluded_mask, a, resid)`` where ``a`` is
        the MoGe->render depth scale actually used this frame (freshly fit, or
        carried over from ``last_a`` when too few core pixels supported a fit)
        and ``resid`` is its residual (``nan`` when carried over). ``drawn`` and
        ``occluded`` depend only on the robot render and the MoGe depth, never
        on ``frame_bgr``/``background_bgr`` -- so they are identical whichever
        background was composited under them; ``render_background_probe.py``
        checks this rather than assuming it.
    """
    rr = rend.render(robot.link_poses(joint_row, float(gripper_value)), camera)
    core = cv2.erode(rr.mask.astype(np.uint8), _ERODE).astype(bool)
    fit = core & np.isfinite(moge) & (rr.depth > 0)
    if fit.sum() >= 200:
        a, resid = robust_scale(moge[fit], rr.depth[fit])
    else:
        a, resid = last_a, float("nan")  # hold the last good scale
    aligned = a * moge
    drawn = rr.mask & (rr.depth > 0)
    occluded = drawn & np.isfinite(aligned) & (rr.depth >= aligned + depth_tol_m)
    out = (frame_bgr if background_bgr is None else background_bgr).copy()
    win = drawn & ~occluded
    out[win] = rr.color[win][:, ::-1]  # RenderResult.color is RGB
    return out, drawn, occluded, a, resid


class ClipBuilder:
    """Holds the URDF and one renderer across every clip of a run.

    Both are expensive to construct (mesh load + EGL context) and neither depends
    on the episode, so they are built once and reused.
    """

    def __init__(self, urdf: Path, table: dict, depth_tol_m: float, crf: int) -> None:
        from fpgm.robot.urdf import RobotModel

        self.robot = RobotModel(str(urdf), load_meshes=True)
        self.meshes = self.robot.visual_meshes()
        self.table = table
        self.depth_tol_m = depth_tol_m
        self.crf = crf
        self._renderer = None
        self._size: tuple[int, int] | None = None

    def renderer(self, w: int, h: int):
        from fpgm.robot.render import RobotRenderer

        if self._size != (w, h):
            if self._renderer is not None:
                self._renderer.close()
            self._renderer = RobotRenderer(self.meshes, w, h)
            self._size = (w, h)
        return self._renderer

    def close(self) -> None:
        if self._renderer is not None:
            self._renderer.close()

    def load_clip(self, depth_meta: dict, episodes_root: Path, cameras_dir: Path) -> dict:
        """Every per-clip input the render loop needs, minus the depth array.

        Split out of :meth:`build` so a caller that wants to render this same
        clip with something other than the recorded gripper trace (see
        ``scripts/render_counterfactual_grip.py``) can get ``joint_positions``,
        the real ``gripper`` trace, the camera and the frame/row alignment
        without duplicating the extrinsics-selection and fps-ratio logic --
        both have already bitten this pipeline once each (see the module and
        ``frames_per_step`` docstrings) and should not get a second copy to
        drift out of sync.
        """
        from fpgm.datagen.robot_buffers import (
            TRAJECTORY_GRIPPER_POSITION_KEY,
            TRAJECTORY_JOINT_POSITIONS_KEY,
            load_extrinsics_candidates,
        )
        from fpgm.geometry.camera import Camera

        uuid, serial = depth_meta["uuid"], depth_meta["camera_serial"]
        ep = episodes_root / uuid
        traj_path = ep / "trajectory.h5"
        mp4 = REPO_ROOT / depth_meta["mp4"]
        src_w, src_h = depth_meta["source_wh"]
        out_w, out_h = depth_meta["video_wh"]
        n = depth_meta["n_frames"]

        with h5py.File(traj_path, "r") as f:
            joint_positions = np.asarray(f[TRAJECTORY_JOINT_POSITIONS_KEY])
            gripper = np.asarray(f[TRAJECTORY_GRIPPER_POSITION_KEY])
        n_rows = joint_positions.shape[0]
        # The mp4-fps trap: the container advertises 60/1 but its frames are 15 Hz
        # trajectory samples. The only trustworthy map is the frame-count ratio.
        cap_probe = cv2.VideoCapture(str(mp4))
        n_video = int(cap_probe.get(cv2.CAP_PROP_FRAME_COUNT)) or n
        cap_probe.release()
        frames_per_step = n_video / n_rows

        intr_src, intr_source = intrinsics_for(self.table, serial, src_w, src_h)
        intr = crop_intrinsics(intr_src, out_w, out_h)
        cands = load_extrinsics_candidates(
            cameras_dir / f"{uuid}_cameras.json", traj_path, serial)
        chosen = next((c for c in cands if c.name == "optimized_cameras_json"), cands[0])
        camera = Camera(intr, chosen.world_to_cam)
        return {
            "uuid": uuid, "serial": serial, "mp4": mp4, "traj_path": traj_path,
            "joint_positions": joint_positions, "gripper": gripper, "n_rows": n_rows,
            "frames_per_step": frames_per_step, "camera": camera,
            "out_w": out_w, "out_h": out_h, "n": n,
            "intr_source": intr_source, "intr_src": intr_src, "extrinsics": chosen.name,
        }

    def render_loop(
        self, joint_positions: np.ndarray, gripper: np.ndarray, mp4: Path, camera,
        frames_per_step: float, n: int, depth_all, out_w: int, out_h: int,
        out_dir: Path, fps: float, *, gripper_reference: np.ndarray | None = None,
    ) -> dict:
        """Per-frame render + MoGe-occlusion composite. Writes ``control.mp4`` /
        ``target.mp4`` into ``out_dir`` and returns the stats :meth:`build` folds
        into ``meta.json``.

        ``gripper_reference``, if given, is a second gripper trace (same length
        and row-indexing as ``gripper``) rendered alongside it purely to measure
        how far the two diverge on screen -- ``render_counterfactual_grip.py``
        passes the real recorded trace here while rendering the overridden one,
        to quantify how much of the real arm/gripper is left showing through the
        composite where the two silhouettes disagree. It costs one extra render
        per frame and changes no output when omitted (the default, and what
        :meth:`build` uses), so this reproduces the loop it was extracted from.
        """
        n_rows = joint_positions.shape[0]
        rend = self.renderer(out_w, out_h)
        stats = {"a": [], "resid": [], "reject": [], "robot_px": []}
        diff = {"extra_real_px": [], "drawn_real_px": []} if gripper_reference is not None else None
        last_a = 1.0
        control_w = H264Writer(out_dir / "control.mp4", out_w, out_h, fps, self.crf)
        target_w = H264Writer(out_dir / "target.mp4", out_w, out_h, fps, self.crf)
        cap = cv2.VideoCapture(str(mp4))
        try:
            for t in range(n):
                ok, frame_bgr = cap.read()
                if not ok:
                    break
                frame_bgr = cover_crop(frame_bgr, out_w, out_h)
                target_w.write(frame_bgr)

                row = min(int(round(t / frames_per_step)), n_rows - 1)
                moge = cv2.resize(depth_all[t].astype(np.float32), (out_w, out_h),
                                  interpolation=cv2.INTER_NEAREST)
                out, drawn, occluded, last_a, resid = render_and_composite(
                    rend, self.robot, joint_positions[row], gripper[row],
                    frame_bgr, moge, camera, self.depth_tol_m, last_a)
                stats["a"].append(last_a)
                stats["resid"].append(resid)
                stats["robot_px"].append(int(drawn.sum()))
                stats["reject"].append(int(occluded.sum()))
                control_w.write(out)

                if diff is not None:
                    rr_ref = rend.render(
                        self.robot.link_poses(joint_positions[row],
                                              float(gripper_reference[row])),
                        camera)
                    drawn_ref = rr_ref.mask & (rr_ref.depth > 0)
                    diff["extra_real_px"].append(int((drawn_ref & ~drawn).sum()))
                    diff["drawn_real_px"].append(int(drawn_ref.sum()))
        finally:
            cap.release()
            control_w.close()
            target_w.close()
        if diff is not None:
            stats["diff"] = diff
        return stats

    def build(self, depth_meta: dict, work: Path, out_root: Path,
              episodes_root: Path, cameras_dir: Path, fps: float) -> dict:
        uuid, serial = depth_meta["uuid"], depth_meta["camera_serial"]
        clip = self.load_clip(depth_meta, episodes_root, cameras_dir)
        joint_positions, gripper = clip["joint_positions"], clip["gripper"]
        mp4, camera = clip["mp4"], clip["camera"]
        frames_per_step, n = clip["frames_per_step"], clip["n"]
        out_w, out_h = clip["out_w"], clip["out_h"]
        intr_src, intr_source = clip["intr_src"], clip["intr_source"]

        depth_all = np.load(work / f"{uuid}__{serial}.depth.npy", mmap_mode="r")
        out_dir = out_root / f"{uuid}__{serial}"
        out_dir.mkdir(parents=True, exist_ok=True)

        stats = self.render_loop(joint_positions, gripper, mp4, camera, frames_per_step,
                                 n, depth_all, out_w, out_h, out_dir, fps)

        ep = episodes_root / uuid
        ep_meta = json.loads((ep / "metadata.json").read_text())
        src_w, src_h = depth_meta["source_wh"]
        a = np.asarray(stats["a"])
        robot_px = max(int(np.sum(stats["robot_px"])), 1)
        meta = {
            "uuid": uuid, "camera_serial": serial, "n_frames": len(stats["a"]),
            "fps": fps, "video_wh": [out_w, out_h], "source_wh": [src_w, src_h],
            "crf": self.crf,
            "caption": ep_meta.get("current_task"), "lab": ep_meta.get("lab"),
            "control": rel(out_dir / "control.mp4"),
            "target": rel(out_dir / "target.mp4"),
            "source_mp4": rel(mp4),
            "intrinsics_source": intr_source,
            "intrinsics_fx_over_w": round(intr_src.fx / src_w, 5),
            "moge_fx_over_w_median": depth_meta["moge_fx_over_w_median"],
            "extrinsics": clip["extrinsics"],
            "frames_per_step": round(frames_per_step, 5),
            "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(stats["resid"])), 4),
            "robot_px_median": int(np.median(stats["robot_px"])),
            "occluded_px_fraction": round(float(np.sum(stats["reject"]) / robot_px), 5),
            "depth_tol_m": self.depth_tol_m,
        }
        (out_dir / "meta.json").write_text(json.dumps(meta, indent=2))
        return meta


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--work", type=Path, required=True)
    ap.add_argument("--out", type=Path, default=REPO_ROOT / "outputs/appearance_pairs")
    ap.add_argument("--episodes-root", type=Path, default=REPO_ROOT / "data/droid_raw")
    ap.add_argument("--cameras-dir", type=Path,
                    default=REPO_ROOT / "data/pointworld/droid/cameras")
    ap.add_argument("--intrinsics", type=Path,
                    default=REPO_ROOT / "configs/droid_camera_intrinsics.json")
    ap.add_argument("--urdf", type=Path, default=None)
    ap.add_argument("--depth-tol-m", type=float, default=0.15)
    ap.add_argument("--crf", type=int, default=18)
    ap.add_argument("--fps", type=float, default=_TRAJECTORY_FPS)
    ap.add_argument("--follow", action="store_true",
                    help="keep consuming until every stage-1 shard drops its "
                         "STAGE1_DONE_* marker")
    ap.add_argument("--expect-shards", type=int, default=1)
    ap.add_argument("--keep-depth", action="store_true",
                    help="do not delete each depth npy after consuming it")
    args = ap.parse_args()

    urdf = args.urdf
    if urdf is None:
        from fpgm.config_datagen import DatagenProfile
        urdf = Path(DatagenProfile.from_yaml(
            str(REPO_ROOT / "configs/datagen_droid.yaml")).paths.urdf)

    builder = ClipBuilder(urdf, load_intrinsic_table(args.intrinsics),
                          args.depth_tol_m, args.crf)
    args.out.mkdir(parents=True, exist_ok=True)
    done, failed, frames, t0 = 0, 0, 0, time.time()
    try:
        while True:
            todo = [p for p in sorted(args.work.glob("*.depth.json"))
                    if p.with_suffix(".npy").exists()
                    and not (args.out / p.name.replace(".depth.json", "")
                             / "meta.json").exists()]
            if not todo:
                if not args.follow or len(list(args.work.glob("STAGE1_DONE_*"))) \
                        >= args.expect_shards:
                    break
                time.sleep(2)   # poll interval: short enough not to show up in
                continue        # the tail latency of a run that is nearly done
            for meta_path in todo:
                dm = json.loads(meta_path.read_text())
                try:
                    m = builder.build(dm, args.work, args.out, args.episodes_root,
                                      args.cameras_dir, args.fps)
                    done += 1
                    frames += m["n_frames"]
                    print(f"[{done}] {m['uuid']}/{m['camera_serial']} "
                          f"n={m['n_frames']} a={m['scale_a_median']:.3f} "
                          f"resid={m['align_residual_m_median']*1000:.1f}mm "
                          f"occl={m['occluded_px_fraction']*100:.2f}% "
                          f"intr={m['intrinsics_source']}", flush=True)
                except Exception as exc:
                    failed += 1
                    print(f"  FAILED {dm['uuid']}/{dm['camera_serial']}: "
                          f"{type(exc).__name__}: {exc}", flush=True)
                if not args.keep_depth:
                    meta_path.with_suffix(".npy").unlink(missing_ok=True)
    finally:
        builder.close()
    dt = time.time() - t0
    print(f"stage 2 done: {done} clips, {failed} failed, {dt/60:.1f} min "
          f"({dt/max(done,1):.1f} s/clip, {frames/max(dt,1e-9):.1f} fps)", flush=True)


if __name__ == "__main__":
    main()