scripts/appearance_depth.py

#!/usr/bin/env python
"""Stage 1 of the appearance-pair pipeline: MoGe depth for the occlusion test.

This is deliberately the *smallest* thing stage 2 cannot compute for itself, and
every choice below is a measured one on an H200 at 1280x720:

* **uint8 -> GPU -> float32**, not ``torch.tensor(rgb / 255.0)``. That innocent
  division promotes a uint8 frame to a **float64** array -- 22 MB per frame,
  built and copied on the CPU -- before torch ever sees it. Removing it takes
  the model from 13.5 to **24.1 fps**, a 1.8x speedup for a one-line change and
  no change in output.
* **Batched inference** (``--batch 4``): 24.1 -> **31.4 fps**. Beyond 4 the
  curve is flat (32.1 at batch 8), so 4 is where the VRAM stops buying speed.
* **Cover-cropped to the training resolution first.** MoGe itself does not care
  -- it resizes to a token budget internally, so 832x480 scores 31.5 fps against
  31.4 at 1280x720 -- but everything downstream of it does: stage 2's composite
  runs 46 fps at 720p against 159 at 832x480. Cropping here means the crop is
  applied exactly once, by one piece of code, to the frame that both stages and
  the trainer see.
* **fp16 depth at half the output resolution**, and nothing else. This is not
  the ``moge_run_*`` cache: those write a 5.6 MB/frame ``.npz`` (depth f32 +
  mask + normal + intrinsics), and that compression -- not the network -- is
  what caps them at 2.45 fps. The only consumer here is a depth *comparison*
  against a URDF render with a 150 mm tolerance, which 416x240 fp16 serves at
  200 KB/frame.

Decoding, at 61-74 fps, is now the co-bottleneck with the model. That is why
stage 1 and stage 2 stay separate processes rather than merging: run
sequentially in one process the whole chain would be ~14 fps, whereas two
processes overlap decode/GPU/encode and both settle near 21 fps.

MoGe's own estimated intrinsics are recorded but NOT used downstream -- stage 2
uses the camera's real calibration (``configs/droid_camera_intrinsics.json``),
because the extrinsics in ``cameras.json`` were optimised against that. MoGe's
estimate is kept as a cheap independent check: on the RAIL demo clip it lands at
fx/w 0.4151 against the true 0.4107.

Env: ``moge``. Usage:
    CUDA_VISIBLE_DEVICES=1 /home/quang/miniconda3/envs/moge/bin/python -u \\
        scripts/appearance_depth.py --work <dir> [--shard 0 3]
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import cv2
import numpy as np
import torch

REPO_ROOT = Path(__file__).resolve().parents[1]

#: VACE's training resolution. Both stages and the trainer work here.
DEFAULT_RESOLUTION = (832, 480)

#: Depth is stored at half the output resolution. The occlusion test tolerates a
#: 150 mm margin, so a depth pixel covering a 2x2 output block costs nothing it
#: can measure, and it turns 800 KB/frame of intermediate into 200 KB.
_DEPTH_SCALE = 0.5


def cover_crop_params(in_w: int, in_h: int, out_w: int, out_h: int):
    """``(resized_w, resized_h, off_x, off_y)`` for a scale-to-cover + centre-crop.

    Same arithmetic as ``fpgm.datagen.export_vace._CoverCrop``, kept here so
    stage 1 (which runs in the ``moge`` env, without ``fpgm`` importable) and
    stage 2 cannot drift apart on the framing.
    """
    scale = max(out_w / in_w, out_h / in_h)
    rw, rh = int(round(in_w * scale)), int(round(in_h * scale))
    return rw, rh, (rw - out_w) / 2.0, (rh - out_h) / 2.0


def cover_crop(image: np.ndarray, out_w: int, out_h: int) -> np.ndarray:
    in_h, in_w = image.shape[:2]
    if (in_w, in_h) == (out_w, out_h):
        return image
    rw, rh, ox, oy = cover_crop_params(in_w, in_h, out_w, out_h)
    interp = cv2.INTER_AREA if rw < in_w else cv2.INTER_LINEAR
    resized = cv2.resize(image, (rw, rh), interpolation=interp)
    x0, y0 = int(round(ox)), int(round(oy))
    return resized[y0:y0 + out_h, x0:x0 + out_w]


def episode_jobs(episodes_root: Path, uuids: list[str] | None) -> list[tuple[str, str, Path]]:
    """``(uuid, camera_serial, mp4)`` for every exterior camera of every episode.

    Both exterior cameras are used: two genuinely different viewpoints of the
    same trajectory, so an episode yields two training clips at no extra download
    cost. The wrist camera is excluded -- the robot fills the frame and the URDF
    render has nothing to composite against.
    """
    jobs: list[tuple[str, str, Path]] = []
    for ep in sorted(p for p in episodes_root.iterdir() if p.is_dir()):
        if uuids and ep.name not in uuids:
            continue
        meta_path = ep / "metadata.json"
        if not meta_path.exists():
            continue
        meta = json.loads(meta_path.read_text())
        for key in ("ext1_cam_serial", "ext2_cam_serial"):
            serial = meta.get(key)
            if not serial:
                continue
            mp4 = ep / "recordings/MP4" / f"{serial}.mp4"
            if mp4.exists():
                jobs.append((ep.name, str(serial), mp4))
    return jobs


def run_clip(model, uuid: str, serial: str, mp4: Path, work: Path,
             device: torch.device, resolution: tuple[int, int],
             resolution_level: int, batch: int) -> dict:
    stem = work / f"{uuid}__{serial}"
    meta_path = stem.with_suffix(".depth.json")
    if meta_path.exists() and stem.with_suffix(".depth.npy").exists():
        return json.loads(meta_path.read_text())

    out_w, out_h = resolution
    cap = cv2.VideoCapture(str(mp4))
    src_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    src_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    dh, dw = int(out_h * _DEPTH_SCALE), int(out_w * _DEPTH_SCALE)

    buf = np.empty((n, dh, dw), np.float16)
    fx_over_w: list[float] = []
    t = 0
    t0 = time.time()
    pending: list[np.ndarray] = []

    def flush() -> None:
        """Run one batch and store its depth. Empties ``pending``."""
        nonlocal t
        if not pending:
            return
        # uint8 -> GPU -> float32 in one hop. Going through `rgb / 255.0` first
        # would materialise a float64 array per frame on the CPU (1.8x slower).
        arr = torch.from_numpy(np.stack(pending)).to(device, non_blocking=True)
        arr = arr.permute(0, 3, 1, 2).float().div_(255.0)
        out = model.infer(arr, resolution_level=resolution_level)
        depth = out["depth"].float().cpu().numpy()
        mask = out["mask"].cpu().numpy().astype(bool)
        k = out["intrinsics"].float().cpu().numpy()
        for i in range(len(pending)):
            # Invalid depth becomes +inf, not NaN: stage 2 asks "is the scene in
            # front of the robot", and inf is the honest answer for "unknown"
            # (it never occludes) whereas NaN would need a separate mask array.
            d = np.where(mask[i], depth[i], np.inf)
            buf[t] = cv2.resize(d, (dw, dh), interpolation=cv2.INTER_NEAREST)
            fx_over_w.append(float(k[i, 0, 0]))
            t += 1
        pending.clear()

    with torch.inference_mode():
        while t + len(pending) < n:
            ok, frame_bgr = cap.read()
            if not ok:
                break
            pending.append(cover_crop(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB),
                                      out_w, out_h))
            if len(pending) == batch:
                flush()
        flush()
    cap.release()
    if t == 0:
        raise RuntimeError(f"{uuid}/{serial}: decoded 0 frames from {mp4}")
    dt = time.time() - t0

    np.save(stem.with_suffix(".depth.npy"), buf[:t])
    meta = {"uuid": uuid, "camera_serial": serial, "mp4": str(mp4.relative_to(REPO_ROOT)),
            "n_frames": t, "source_wh": [src_w, src_h], "video_wh": [out_w, out_h],
            "depth_wh": [dw, dh],
            "moge_fx_over_w_median": round(float(np.median(fx_over_w)), 5),
            "seconds": round(dt, 1), "fps": round(t / dt, 2)}
    meta_path.write_text(json.dumps(meta))
    print(f"  {uuid}/{serial}: {t} frames {dt:.1f}s ({t/dt:.2f} fps) "
          f"moge fx/w={meta['moge_fx_over_w_median']:.4f}", flush=True)
    return meta


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--episodes-root", type=Path, default=REPO_ROOT / "data/droid_raw")
    ap.add_argument("--work", type=Path, required=True, help="scratch dir for depth npy")
    ap.add_argument("--model", default="Ruicheng/moge-2-vitl-normal")
    ap.add_argument("--resolution", type=int, nargs=2, default=DEFAULT_RESOLUTION,
                    metavar=("W", "H"))
    ap.add_argument("--resolution-level", type=int, default=9)
    ap.add_argument("--batch", type=int, default=4,
                    help="frames per forward pass; 4 captures the whole speedup "
                         "(31.4 fps vs 32.1 at 8, against 24.1 at 1)")
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--uuid", nargs="*", default=None)
    ap.add_argument("--uuid-file", type=Path, default=None,
                    help="newline-separated uuids, e.g. from sample_appearance_episodes.py")
    ap.add_argument("--shard", type=int, nargs=2, default=(0, 1), metavar=("I", "N"),
                    help="process clips i, i+N, i+2N... Round-robin, so long clips "
                         "spread evenly when running one shard per GPU.")
    ap.add_argument("--max-pending", type=int, default=16,
                    help="pause while this many un-consumed depth files are queued, "
                         "so the scratch dir cannot outrun stage 2 and fill the disk")
    args = ap.parse_args()

    uuids = list(args.uuid or [])
    if args.uuid_file:
        uuids += [u for u in args.uuid_file.read_text().split() if u]
    jobs = episode_jobs(args.episodes_root, uuids or None)
    i, nsh = args.shard
    jobs = jobs[i::nsh]
    if args.limit:
        jobs = jobs[: args.limit]
    args.work.mkdir(parents=True, exist_ok=True)
    print(f"shard {i}/{nsh}: {len(jobs)} clips from {args.episodes_root} "
          f"@ {args.resolution[0]}x{args.resolution[1]} batch={args.batch}", flush=True)

    from moge.model.v2 import MoGeModel
    device = torch.device("cuda")
    model = MoGeModel.from_pretrained(args.model).to(device).eval()

    for j, (uuid, serial, mp4) in enumerate(jobs):
        while len(list(args.work.glob("*.depth.npy"))) >= args.max_pending:
            time.sleep(5)
        try:
            run_clip(model, uuid, serial, mp4, args.work, device,
                     tuple(args.resolution), args.resolution_level, args.batch)
        except Exception as exc:  # one bad clip must not end a multi-hour run
            print(f"  FAILED {uuid}/{serial}: {type(exc).__name__}: {exc}", flush=True)
        if (j + 1) % 20 == 0:
            print(f"[{j+1}/{len(jobs)}]", flush=True)
    (args.work / f"STAGE1_DONE_{i}of{nsh}").write_text("")
    print("stage 1 done", flush=True)


if __name__ == "__main__":
    main()