scripts/moge_render_panels.py

#!/usr/bin/env python
"""Re-render the MoGe preview videos from the cached ``per_frame/*.npz`` — no GPU.

``moge_run_clips.py`` writes a three-panel RGB|depth|normal video, and on this
project's clips that video *understates* the model: the global depth range is
~1.1 m spread over 255 colour levels, i.e. **4.3 mm per level**, so the 17 mm step
of the manipulated LEGO brick occupies about four levels of turbo and reads as flat
by eye. It is not flat -- measured at frame 90 of the demo episode, the mean depth
inside the brick mask sits 17.3 mm in front of the surrounding ring.

So this adds a fourth panel: the **same depth, near-field colour range**, clipped to
the clip's [p25, p60] depth percentiles -- roughly the manipulation working volume,
~1 mm per colour level, where that 17 mm step spans ~16 levels and is obvious.

This is deliberately NOT a high-pass / unsharp filter. A ``depth - blur(depth)``
panel was tried first and rejected: with any sigma, real depth discontinuities
(arm against background) dwarf object relief -- the p99 of |high-pass| is ~520 mm
against the ~17 mm being looked for -- so the panel saturates into blobs and stops
being readable as depth at all. A narrower colour window applies no filter, keeps
every pixel's value literal, and prints its own mm-per-level in the label, so what
is on screen can be checked against a number.

Layout is 2x2 (2560x1440) rather than 1x4: a 5120-px-wide video gets downscaled by
every player and the detail this exists to show is the first thing lost.

Env: anything with numpy/cv2 (``fpgm`` is fine); imports no torch and no ``fpgm``.

Usage:
    /home/quang/miniconda3/envs/fpgm/bin/python scripts/moge_render_panels.py \\
        [--clips-root outputs/moge/published_clips] [--near-percentiles 25 60]
"""

from __future__ import annotations

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

import cv2
import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
_FFMPEG_CANDIDATE = Path("/home/quang/miniconda3/envs/ffmpeg_libs/bin/ffmpeg")
_PERCENTILES = (2.0, 98.0)
#: See moge_run_clips.TRAJECTORY_FPS -- DROID containers say 60/1, content is 15 Hz.
TRAJECTORY_FPS = 15.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("no ffmpeg found")


def write_h264(frames_bgr, out: Path, fps: float, crf: int = 16) -> None:
    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: np.ndarray, text: str) -> np.ndarray:
    out = img.copy()
    for colour, thick in ((0, 0, 0), 5), ((255, 255, 255), 2):
        cv2.putText(out, text, (14, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.95, colour,
                    thick, cv2.LINE_AA)
    return out


def colorize_depth(depth: np.ndarray, lo: float, hi: float) -> np.ndarray:
    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)
    col = cv2.applyColorMap(((1.0 - norm) * 255).astype(np.uint8), cv2.COLORMAP_TURBO)
    col[~valid] = 0
    return col


def colorize_normal(normal: np.ndarray, valid: np.ndarray) -> np.ndarray:
    bgr = (((normal.astype(np.float32) * 0.5 + 0.5) * 255).clip(0, 255)
           .astype(np.uint8))[..., ::-1].copy()
    bgr[~valid] = 0
    return bgr


def render_clip(clip_dir: Path, near_pcts: tuple[float, float], fps: float) -> dict:
    npz_dir = clip_dir / "per_frame"
    files = sorted(npz_dir.glob("frame_*.npz"))
    if not files:
        raise FileNotFoundError(f"no npz under {npz_dir}")
    stats = json.loads((clip_dir / "stats.json").read_text())
    video = REPO_ROOT / stats["video"]

    cap = cv2.VideoCapture(str(video))
    rgbs = []
    for _ in range(len(files)):
        ok, fr = cap.read()
        if not ok:
            break
        rgbs.append(fr)
    cap.release()
    n = min(len(rgbs), len(files))
    print(f"  {clip_dir.name}: {n} frames")

    depths, masks, normals = [], [], []
    for f in files[:n]:
        d = np.load(f)
        depths.append(d["depth"])
        masks.append(d["mask"].astype(bool))
        normals.append(d["normal"].astype(np.float32) if "normal" in d else None)

    pool = np.concatenate([d[np.isfinite(d)].ravel()[::17] for d in depths])
    lo, hi = np.percentile(pool, _PERCENTILES)
    nlo, nhi = np.percentile(pool, near_pcts)
    mmpl, near_mmpl = (hi - lo) * 1000 / 255, (nhi - nlo) * 1000 / 255
    print(f"    depth global {lo:.3f}-{hi:.3f} m ({mmpl:.1f} mm/level) | "
          f"near-field p{near_pcts[0]:g}-p{near_pcts[1]:g} = {nlo:.3f}-{nhi:.3f} m "
          f"({near_mmpl:.1f} mm/level)")

    def frames():
        for i in range(n):
            valid = masks[i] & np.isfinite(depths[i])
            tl = label(rgbs[i], "RGB")
            tr = label(colorize_depth(depths[i], lo, hi),
                       f"MoGe-2 depth   {lo:.2f}-{hi:.2f} m  ({mmpl:.1f} mm/level)")
            bl = label(colorize_depth(depths[i], nlo, nhi),
                       f"same depth, near-field range   {nlo:.2f}-{nhi:.2f} m  "
                       f"({near_mmpl:.1f} mm/level)")
            br = label(colorize_normal(normals[i], valid) if normals[i] is not None
                       else np.zeros_like(rgbs[i]), "MoGe-2 normal")
            yield np.vstack([np.hstack([tl, tr]), np.hstack([bl, br])])

    out = clip_dir / "panels_2x2.mp4"
    write_h264(frames(), out, fps)
    print(f"    wrote {out}")
    stats.update({"panels_video": str(out.relative_to(REPO_ROOT)),
                  "near_field_percentiles": list(near_pcts),
                  "near_field_range_m": [float(nlo), float(nhi)],
                  "near_field_mm_per_colour_level": round(near_mmpl, 2),
                  "depth_mm_per_colour_level": round(mmpl, 2),
                  "output_fps": fps})
    (clip_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("--clips-root", type=Path,
                    default=REPO_ROOT / "outputs/moge/published_clips")
    ap.add_argument("--near-percentiles", type=float, nargs=2, default=(25.0, 60.0),
                    metavar=("LO", "HI"),
                    help="depth percentiles bounding the near-field panel's colour range")
    ap.add_argument("--fps", type=float, default=TRAJECTORY_FPS)
    args = ap.parse_args()

    out = []
    for d in sorted(p for p in args.clips_root.iterdir() if p.is_dir()):
        if not (d / "stats.json").exists():
            print(f"  skip {d.name}: no stats.json")
            continue
        out.append(render_clip(d, tuple(args.near_percentiles), args.fps))
    print(f"\nrendered {len(out)} clip(s)")


if __name__ == "__main__":
    main()