scripts/moge_side_by_side.py

#!/usr/bin/env python
"""RGB | MoGe depth side-by-side video for a whole clip.

Run with the ``fpgm`` env (it only needs numpy/cv2/imageio to read the ``.npz``
files ``moge_infer_episode.py`` already wrote -- no torch, no MoGe import).

**One global depth normalisation for the whole clip, never per frame.** Per-frame
normalisation makes the depth panel flicker as the scene's near/far extremes move,
which is indistinguishable, by eye, from the model's own temporal instability --
exactly the thing this video is meant to let someone judge. Robust percentiles, so
one stray pixel cannot crush the range. The measured per-frame scale drift for
MoGe-2 on this project's demo episode was CV 9.6%, so the distinction matters here.

Usage:
    PYTHONPATH=src python scripts/moge_side_by_side.py \\
        --video <clip.mp4> --npz-dir outputs/moge/marker_per_frame \\
        --out outputs/moge/marker_side_by_side.mp4
"""

from __future__ import annotations

import argparse
import subprocess
from pathlib import Path

import cv2
import imageio.v3 as iio
import numpy as np

from fpgm.viz.video import _resolve_ffmpeg

_PERCENTILES = (2.0, 98.0)


def write_h264(frames_bgr, out: Path, fps: float) -> None:
    t, h, w, _ = frames_bgr.shape
    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", "18", str(out),
    ]
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE)
    assert p.stdin is not None
    try:
        for i in range(t):
            p.stdin.write(np.ascontiguousarray(frames_bgr[i]).tobytes())
    finally:
        p.stdin.close()
        if p.wait() != 0:
            raise RuntimeError(f"ffmpeg failed writing {out}")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--video", required=True, type=Path)
    ap.add_argument("--npz-dir", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--fps", type=float, default=15.0)
    ap.add_argument("--panel-width", type=int, default=640)
    args = ap.parse_args()

    npzs = sorted(args.npz_dir.glob("*.npz"))
    if not npzs:
        raise SystemExit(f"no .npz under {args.npz_dir}")
    rgb = np.asarray(iio.imread(args.video, plugin="pyav"))
    n = min(len(rgb), len(npzs))
    print(f"rgb {rgb.shape} | npz {len(npzs)} -> using {n} frames")

    # Pass 1: global range over a subsample (loading 992 full-res depth maps twice
    # is wasteful; every 5th frame is plenty to fix a robust percentile range).
    vals = []
    for p in npzs[: n : max(1, n // 200)]:
        d = np.load(p)
        depth, mask = d["depth"], d["mask"]
        v = depth[mask & np.isfinite(depth) & (depth > 0)]
        if v.size:
            vals.append(np.random.default_rng(0).choice(v, size=min(v.size, 20000), replace=False))
    allv = np.concatenate(vals)
    lo, hi = np.percentile(1.0 / allv, _PERCENTILES)
    print(f"global inverse-depth range: {lo:.4f} .. {hi:.4f} per-m "
          f"(metric {1/hi:.3f} .. {1/lo:.3f} m)")

    pw = args.panel_width
    ph = int(round(pw * rgb.shape[1] / rgb.shape[2]))
    out = np.empty((n, ph, pw * 2, 3), dtype=np.uint8)
    for i in range(n):
        d = np.load(npzs[i])
        depth, mask = d["depth"], d["mask"]
        g = np.zeros(depth.shape, np.float32)
        m = mask & np.isfinite(depth) & (depth > 0)
        if m.any():
            g[m] = np.clip((1.0 / depth[m] - lo) / (hi - lo), 0.0, 1.0)
        col = cv2.applyColorMap((g * 255).astype(np.uint8), cv2.COLORMAP_TURBO)
        col[~m] = 0
        left = cv2.resize(rgb[i][..., ::-1], (pw, ph))
        right = cv2.resize(col, (pw, ph), interpolation=cv2.INTER_NEAREST)
        for panel, text in ((left, f"RGB  f{i:04d}"), (right, "MoGe-2 depth")):
            cv2.rectangle(panel, (0, 0), (len(text) * 11 + 10, 24), (0, 0, 0), -1)
            cv2.putText(panel, text, (5, 17), cv2.FONT_HERSHEY_SIMPLEX, 0.55,
                        (255, 255, 255), 1, cv2.LINE_AA)
        out[i] = np.concatenate([left, right], axis=1)

    write_h264(out, args.out, args.fps)
    print(f"wrote {args.out}  ({n} frames, {pw*2}x{ph})")


if __name__ == "__main__":
    main()