scripts/moge_run_clips.py

#!/usr/bin/env python
"""Run MoGe-2 over the FULL length of several DROID clips and render what it sees.

Purpose is qualitative: "how fine-grained is the monocular depth model on our own
published clips?". So the primary output is a three-panel video per clip

    RGB  |  depth (turbo, ONE global normalisation for the whole clip)  |  normal

rather than a metric table -- normals are included because they are the panel that
actually exposes fine detail: depth is dominated by the metre-scale front-to-back
ramp of the table, while the normal map shows the millimetre-scale surface relief
(the drawer lip, the brick's edges, the cloth folds) on the same colour budget.

Two choices worth stating, because both are the difference between judging the
model and judging the visualisation:

* **Global depth normalisation, never per frame.** Per-frame percentiles make the
  panel breathe as the near/far extremes move, which by eye is indistinguishable
  from the model's own temporal instability -- the exact thing this video exists to
  let someone judge. Percentiles are robust (2/98) over every valid pixel of every
  frame, so one stray pixel cannot crush the range.
* **Out-of-the-box path: no ``fov_x`` is passed**, so MoGe estimates its own FOV
  per frame even though this project knows the true DROID intrinsics. Priming it
  with the calibration would measure a best case we would not get on arbitrary
  video. ``resolution_level=9`` is the default and already maps to the top of the
  model's ``num_tokens_range`` (3600), i.e. the finest setting it offers.

Env: ``moge`` (``/home/quang/miniconda3/envs/moge/bin/python``), NOT ``fpgm``.
Deliberately imports nothing from ``fpgm`` so it needs only torch/cv2/numpy.

Usage:
    CUDA_VISIBLE_DEVICES=1 /home/quang/miniconda3/envs/moge/bin/python \\
        scripts/moge_run_clips.py --manifest outputs/datagen/_publish_staging/dataset_manifest.json \\
        --limit 5 --out-root outputs/moge/published_clips
"""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import time
from pathlib import Path

import cv2
import numpy as np
import torch

REPO_ROOT = Path(__file__).resolve().parents[1]
_FFMPEG_CANDIDATE = Path("/home/quang/miniconda3/envs/ffmpeg_libs/bin/ffmpeg")
_PERCENTILES = (2.0, 98.0)


def resolve_ffmpeg() -> str:
    if _FFMPEG_CANDIDATE.exists():
        return str(_FFMPEG_CANDIDATE)
    on_path = shutil.which("ffmpeg")
    if on_path:
        return on_path
    raise RuntimeError(f"no ffmpeg: {_FFMPEG_CANDIDATE} missing and none on PATH")


def write_h264(frames_bgr, out: Path, fps: float, crf: int = 16) -> None:
    """Pipe BGR frames to libx264. `frames_bgr` is any iterable of (H, W, 3) uint8."""
    it = iter(frames_bgr)
    first = next(it)
    h, w = first.shape[:2]
    out.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        resolve_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),
    ]
    p = subprocess.Popen(cmd, 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 label(img_bgr: np.ndarray, text: str) -> np.ndarray:
    out = img_bgr.copy()
    cv2.putText(out, text, (14, 34), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 0), 5, cv2.LINE_AA)
    cv2.putText(out, text, (14, 34), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2, cv2.LINE_AA)
    return out


def colorize_depth(depth: np.ndarray, lo: float, hi: float) -> np.ndarray:
    """(H, W) metres -> BGR turbo. Invalid (inf/nan) pixels stay pure black."""
    valid = np.isfinite(depth)
    norm = np.zeros_like(depth, dtype=np.float32)
    norm[valid] = np.clip((depth[valid] - lo) / max(hi - lo, 1e-6), 0.0, 1.0)
    # near = warm: invert so the closest surfaces are the brightest, matching how the
    # RGB panel reads (the manipulated object is always the near thing).
    u8 = ((1.0 - norm) * 255.0).astype(np.uint8)
    col = cv2.applyColorMap(u8, cv2.COLORMAP_TURBO)
    col[~valid] = 0
    return col


def colorize_normal(normal: np.ndarray, valid: np.ndarray) -> np.ndarray:
    """(H, W, 3) unit normals in camera frame -> BGR. Invalid pixels black."""
    rgb = ((normal * 0.5 + 0.5) * 255.0).clip(0, 255).astype(np.uint8)
    bgr = rgb[..., ::-1].copy()
    bgr[~valid] = 0
    return bgr


def select_clips(manifest: Path, limit: int, explicit: list[str] | None) -> list[dict]:
    data = json.loads(manifest.read_text())
    eps = data["episodes"]
    if explicit:
        wanted = set(explicit)
        eps = [e for e in eps if e["uuid"] in wanted or f"{e['uuid']}/{e['camera_serial']}" in wanted]
    # Prefer distinct viewpoints: keep at most one entry per (uuid, serial), and
    # order so the alternate camera of a repeated uuid is not dropped by --limit.
    seen: set[tuple[str, str]] = set()
    uniq = []
    for e in eps:
        key = (e["uuid"], e["camera_serial"])
        if key in seen:
            continue
        seen.add(key)
        uniq.append(e)
    return uniq[:limit] if limit else uniq


#: DROID mp4 containers advertise 60/1, but their frames are the 15 Hz trajectory
#: samples -- the whole repo indexes poses by video frame at
#: ``fpgm.data.episode.DEFAULT_TRAJECTORY_FPS = 15.0`` (e.g. this episode's 178
#: frames carry 178 pose-provenance entries). Muxing the panel video at the
#: container's 60 would play it 4x fast, which reads as temporal jitter that isn't
#: there. Hardcoding the repo constant instead, and recording both in stats.json.
TRAJECTORY_FPS = 15.0


def run_clip(model, ep: dict, out_root: Path, device: torch.device,
             save_npz: bool, resolution_level: int, out_fps: float,
             skip_existing: bool) -> dict:
    uuid, serial = ep["uuid"], ep["camera_serial"]
    video = REPO_ROOT / "data" / "droid_raw" / uuid / "recordings" / "MP4" / f"{serial}.mp4"
    if not video.exists():
        raise FileNotFoundError(video)

    out_dir = out_root / f"{uuid}__{serial}"
    npz_dir = out_dir / "per_frame"
    done_marker = out_dir / "stats.json"
    if skip_existing and done_marker.exists() and (out_dir / "rgb_depth_normal.mp4").exists():
        print(f"\n=== {uuid} / {serial} :: already done, skipping (delete {done_marker} to redo)")
        return json.loads(done_marker.read_text())
    if save_npz:
        npz_dir.mkdir(parents=True, exist_ok=True)
    out_dir.mkdir(parents=True, exist_ok=True)

    cap = cv2.VideoCapture(str(video))
    n_reported = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    container_fps = float(cap.get(cv2.CAP_PROP_FPS)) or out_fps
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    print(f"\n=== {uuid} / {serial} :: {w}x{h}, {n_reported} frames reported, "
          f"container fps {container_fps:.2f} -> playing at {out_fps:.2f}")

    rgbs: list[np.ndarray] = []      # uint8 BGR, for the left panel
    depths: list[np.ndarray] = []    # float16 metres, inf where invalid
    normals: list[np.ndarray] = []   # uint8 BGR, already colourised (keeps RAM flat)

    t0 = time.time()
    with torch.inference_mode():
        for t in range(n_reported):
            ok, frame_bgr = cap.read()
            if not ok:
                print(f"  decode stopped at t={t} (container reported {n_reported})")
                break
            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            image = torch.tensor(frame_rgb / 255.0, dtype=torch.float32,
                                 device=device).permute(2, 0, 1)
            out = model.infer(image, resolution_level=resolution_level)  # no fov_x on purpose

            depth = out["depth"].float().cpu().numpy()
            mask = out["mask"].cpu().numpy().astype(bool)
            normal = out["normal"].float().cpu().numpy() if "normal" in out else None
            intr = out["intrinsics"].float().cpu().numpy()

            if save_npz:
                payload = {"depth": depth.astype(np.float32), "mask": mask,
                           "intrinsics": intr.astype(np.float32)}
                if normal is not None:
                    payload["normal"] = normal.astype(np.float16)
                np.savez_compressed(npz_dir / f"frame_{t:04d}.npz", **payload)

            rgbs.append(frame_bgr)
            depths.append(depth.astype(np.float16))
            normals.append(colorize_normal(normal, mask) if normal is not None
                           else np.zeros_like(frame_bgr))
    cap.release()
    n = len(rgbs)
    dt = time.time() - t0
    print(f"  {n} frames in {dt:.1f}s ({n / max(dt, 1e-6):.2f} fps)")

    # --- one global depth range for the whole clip -------------------------------
    stacked = np.concatenate([d[np.isfinite(d)].astype(np.float32).ravel() for d in depths])
    lo, hi = np.percentile(stacked, _PERCENTILES)
    valid_frac = float(stacked.size) / max(n * h * w, 1)
    print(f"  depth p2={lo:.3f} m  p98={hi:.3f} m  valid_pixels={valid_frac * 100:.1f}%")

    def panels():
        for i in range(n):
            d = colorize_depth(depths[i].astype(np.float32), lo, hi)
            yield np.hstack([
                label(rgbs[i], "RGB"),
                label(d, f"MoGe-2 depth  [{lo:.2f}-{hi:.2f} m, global]"),
                label(normals[i], "MoGe-2 normal"),
            ])

    vid_out = out_dir / "rgb_depth_normal.mp4"
    write_h264(panels(), vid_out, out_fps)
    print(f"  wrote {vid_out}")

    stats = {
        "uuid": uuid, "camera_serial": serial, "camera_role": ep.get("camera_role"),
        "task": ep.get("task"), "video": str(video.relative_to(REPO_ROOT)),
        "width": w, "height": h,
        "container_fps": container_fps, "output_fps": out_fps,
        "n_frames_reported": n_reported, "n_frames_processed": n,
        "seconds": round(dt, 2), "fps_inference": round(n / max(dt, 1e-6), 3),
        "depth_p2_m": float(lo), "depth_p98_m": float(hi),
        "valid_pixel_fraction": round(valid_frac, 5),
        "resolution_level": resolution_level,
        "video_out": str(vid_out.relative_to(REPO_ROOT)),
    }
    (out_dir / "stats.json").write_text(json.dumps(stats, indent=2))
    return stats


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--manifest", type=Path,
                    default=REPO_ROOT / "outputs/datagen/_publish_staging/dataset_manifest.json")
    ap.add_argument("--uuid", nargs="*", default=None,
                    help="explicit uuid(s) or 'uuid/serial'; default: first --limit of the manifest")
    ap.add_argument("--limit", type=int, default=5)
    ap.add_argument("--out-root", type=Path, default=REPO_ROOT / "outputs/moge/published_clips")
    ap.add_argument("--model", default="Ruicheng/moge-2-vitl-normal")
    ap.add_argument("--resolution-level", type=int, default=9,
                    help="9 = default = top of num_tokens_range (3600), the finest MoGe-2 offers")
    ap.add_argument("--no-npz", action="store_true", help="video only, skip the per-frame arrays")
    ap.add_argument("--skip-existing", action=argparse.BooleanOptionalAction, default=True,
                    help="skip clips that already have stats.json + video (default: on)")
    ap.add_argument("--fps", type=float, default=TRAJECTORY_FPS,
                    help="playback fps of the panel video; NOT the mp4 container's 60 (see TRAJECTORY_FPS)")
    args = ap.parse_args()

    clips = select_clips(args.manifest, args.limit, args.uuid)
    print(f"{len(clips)} clip(s) selected:")
    for c in clips:
        print(f"  {c['uuid']} / {c['camera_serial']} ({c.get('camera_role')})")

    from moge.model.v2 import MoGeModel
    device = torch.device("cuda")
    print(f"\nloading {args.model} ...")
    model = MoGeModel.from_pretrained(args.model).to(device).eval()

    all_stats = []
    for ep in clips:
        try:
            all_stats.append(run_clip(model, ep, args.out_root, device,
                                      save_npz=not args.no_npz,
                                      resolution_level=args.resolution_level,
                                      out_fps=args.fps,
                                      skip_existing=args.skip_existing))
        except Exception as exc:  # keep going: one bad clip must not lose the batch
            print(f"  FAILED {ep['uuid']}/{ep['camera_serial']}: {type(exc).__name__}: {exc}")
            all_stats.append({"uuid": ep["uuid"], "camera_serial": ep["camera_serial"],
                              "error": f"{type(exc).__name__}: {exc}"})

    args.out_root.mkdir(parents=True, exist_ok=True)
    summary = args.out_root / "summary.json"
    summary.write_text(json.dumps({"model": args.model,
                                   "resolution_level": args.resolution_level,
                                   "clips": all_stats}, indent=2))
    print(f"\nwrote {summary}")


if __name__ == "__main__":
    main()