scripts/moge_run_sample_bucket.py

#!/usr/bin/env python
"""MoGe-2 depth over every clip of the ``twanghcmut/droid-sample`` bucket.

Unlike ``moge_run_clips.py`` (which serves the 10 processed episodes of the
datagen bucket), this targets the 56-clip diversity sample: 10 labs, 9 object
classes, 23.7k frames, clips from 117 to 2449 frames. That size forces three
choices, each of which was measured rather than assumed:

1. **No per-frame ``.npz``.** At ~5.6 MB/frame compressed the cache would be
   ~133 GB, on a volume sitting at 99%. The panel video is the deliverable here;
   re-running inference is 2.7 GPU-hours, which is cheaper than keeping 133 GB
   around. ``moge_run_clips.py`` still caches, because 10 clips is 5 GB.

2. **Depth kept in RAM, not spilled to disk.** fp16 depth (1.8 MB/frame) plus the
   already-colourised normal panel (2.8 MB/frame) is 4.6 MB/frame -- 11 GB for the
   longest clip, against ~580 GB free. RGB is NOT buffered; the video is simply
   decoded a second time for the render pass, which costs nothing next to the
   forward pass. Spilling instead would have written ~109 GB through an array that
   a concurrent training job is already I/O-bound on.

3. **Global depth range, computed after the full pass -- NOT estimated from a
   prefix.** Estimating the colour range from the first 30 frames was tried and
   rejected on measured evidence: across the 10 datagen clips it moved the p2
   bound by up to +154 mm and the near-field lower bound by up to +106 mm, because
   early frames predate the arm coming close to the camera, so the near extreme is
   simply absent from the prefix. A range that wrong makes the near-field panel
   misleading for most of the clip, which defeats its whole purpose.

Env: ``moge``. Usage:
    CUDA_VISIBLE_DEVICES=2 /home/quang/miniconda3/envs/moge/bin/python -u \\
        scripts/moge_run_sample_bucket.py --index <index.csv> [--shard 0 2]
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
import time
from pathlib import Path

import cv2
import numpy as np
import torch

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "scripts"))
from moge_render_panels import (  # noqa: E402
    TRAJECTORY_FPS, colorize_depth, colorize_normal, label, write_h264,
)

_PERCENTILES = (2.0, 98.0)


def label2(img: np.ndarray, line1: str, line2: str) -> np.ndarray:
    out = label(img, line1)
    for colour, thick in ((0, 0, 0), 4), ((235, 235, 235), 1):
        cv2.putText(out, line2, (14, 68), cv2.FONT_HERSHEY_SIMPLEX, 0.62, colour,
                    thick, cv2.LINE_AA)
    return out


def resolve_video(uuid: str, videos_root: Path) -> Path:
    hits = sorted(videos_root.glob(f"{uuid}/recordings/MP4/*.mp4"))
    if len(hits) != 1:
        raise FileNotFoundError(f"{uuid}: expected exactly 1 mp4, found {len(hits)}")
    return hits[0]


def run_clip(model, row: dict, videos_root: Path, out_root: Path,
             device: torch.device, near_pcts: tuple[float, float],
             resolution_level: int, fps: float, skip_existing: bool,
             save_npz: bool) -> dict:
    uuid = row["uuid"]
    out_dir = out_root / uuid
    npz_dir = out_dir / "per_frame"
    stats_p = out_dir / "stats.json"
    if skip_existing and stats_p.exists() and (out_dir / "panels_2x2.mp4").exists():
        print(f"=== {uuid} :: already done, skipping")
        return json.loads(stats_p.read_text())

    video = resolve_video(uuid, videos_root)
    out_dir.mkdir(parents=True, exist_ok=True)
    if save_npz:
        npz_dir.mkdir(parents=True, exist_ok=True)

    cap = cv2.VideoCapture(str(video))
    n_reported = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    n_index = int(row.get("frames") or 0)
    print(f"=== {uuid} :: {w}x{h}, {n_reported} frames "
          f"({'matches' if n_reported == n_index else f'INDEX SAYS {n_index}'}) "
          f"| {row.get('lab')} / {row.get('object_class')}")

    depths: list[np.ndarray] = []
    normals: list[np.ndarray] = []
    t = 0
    t0 = time.time()
    with torch.inference_mode():
        while True:
            ok, frame_bgr = cap.read()
            if not ok:
                break
            rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            image = torch.tensor(rgb / 255.0, dtype=torch.float32,
                                 device=device).permute(2, 0, 1)
            out = model.infer(image, resolution_level=resolution_level)  # no fov_x
            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
            if save_npz:
                payload = {"depth": depth.astype(np.float32), "mask": mask,
                           "intrinsics": out["intrinsics"].float().cpu().numpy()}
                if normal is not None:
                    payload["normal"] = normal.astype(np.float16)
                np.savez_compressed(npz_dir / f"frame_{t:05d}.npz", **payload)
            depths.append(depth.astype(np.float16))
            normals.append(colorize_normal(normal, mask) if normal is not None
                           else np.zeros((h, w, 3), np.uint8))
            t += 1
    cap.release()
    n = len(depths)
    dt = time.time() - t0
    if n == 0:
        raise RuntimeError(f"{uuid}: decoded 0 frames from {video}")

    pool = np.concatenate([d[np.isfinite(d)].astype(np.float32).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"  {n} frames in {dt:.1f}s ({n / dt:.2f} fps) | global {lo:.3f}-{hi:.3f} m "
          f"({mmpl:.1f} mm/lvl) | near {nlo:.3f}-{nhi:.3f} m ({near_mmpl:.1f} mm/lvl)")

    instr = (row.get("instruction_1") or row.get("current_task") or "")[:78]

    def frames():
        cap2 = cv2.VideoCapture(str(video))
        try:
            for i in range(n):
                ok, rgb_bgr = cap2.read()
                if not ok:
                    break
                d32 = depths[i].astype(np.float32)
                valid = np.isfinite(d32)
                tl = label2(rgb_bgr, f"RGB   {row.get('lab')} / {row.get('object_class')}",
                            instr)
                tr = label(colorize_depth(d32, lo, hi),
                           f"MoGe-2 depth   {lo:.2f}-{hi:.2f} m  ({mmpl:.1f} mm/level)")
                bl = label(colorize_depth(d32, nlo, nhi),
                           f"same depth, near-field range   {nlo:.2f}-{nhi:.2f} m  "
                           f"({near_mmpl:.1f} mm/level)")
                br = label(colorize_normal(np.zeros((h, w, 3), np.float32), valid)
                           if normals[i] is None else normals[i], "MoGe-2 normal")
                yield np.vstack([np.hstack([tl, tr]), np.hstack([bl, br])])
        finally:
            cap2.release()

    out_video = out_dir / "panels_2x2.mp4"
    write_h264(frames(), out_video, fps)
    depths.clear()
    normals.clear()

    stats = {
        "uuid": uuid, "lab": row.get("lab"), "object_class": row.get("object_class"),
        "task_class": row.get("task_class"), "instruction": row.get("instruction_1"),
        "video": str(video.relative_to(REPO_ROOT)),
        "width": w, "height": h, "n_frames_processed": n,
        "n_frames_index_csv": n_index, "output_fps": fps,
        "seconds": round(dt, 2), "fps_inference": round(n / dt, 3),
        "depth_range_m": [float(lo), float(hi)],
        "depth_mm_per_colour_level": round(mmpl, 2),
        "near_field_percentiles": list(near_pcts),
        "near_field_range_m": [float(nlo), float(nhi)],
        "near_field_mm_per_colour_level": round(near_mmpl, 2),
        "resolution_level": resolution_level,
        "panels_video": str(out_video.relative_to(REPO_ROOT)),
        "preview_bytes": out_video.stat().st_size,
    }
    stats_p.write_text(json.dumps(stats, indent=2))
    return stats


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--index", type=Path, required=True, help="the bucket's index.csv")
    ap.add_argument("--videos-root", type=Path, default=REPO_ROOT / "data/droid_coverage56")
    ap.add_argument("--out-root", type=Path, default=REPO_ROOT / "outputs/moge/sample_bucket")
    ap.add_argument("--model", default="Ruicheng/moge-2-vitl-normal")
    ap.add_argument("--resolution-level", type=int, default=9)
    ap.add_argument("--near-percentiles", type=float, nargs=2, default=(25.0, 60.0))
    ap.add_argument("--fps", type=float, default=TRAJECTORY_FPS)
    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 across shards instead of piling into one.")
    ap.add_argument("--skip-existing", action=argparse.BooleanOptionalAction, default=True)
    ap.add_argument("--save-npz", action=argparse.BooleanOptionalAction, default=True,
                    help="keep the full per-frame arrays (default: on)")
    ap.add_argument("--uuid", nargs="*", default=None,
                    help="run only these uuid(s); default: every row of the index")
    args = ap.parse_args()

    rows = list(csv.DictReader(args.index.open()))
    if args.uuid:
        wanted = set(args.uuid)
        found = {r["uuid"] for r in rows} & wanted
        if wanted - found:
            raise SystemExit(f"uuid(s) not in index: {sorted(wanted - found)}")
        rows = [r for r in rows if r["uuid"] in wanted]
    i, nsh = args.shard
    mine = rows[i::nsh]
    print(f"shard {i}/{nsh}: {len(mine)} of {len(rows)} clips, "
          f"{sum(int(r.get('frames') or 0) for r in mine)} frames")

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

    all_stats = []
    for row in mine:
        try:
            all_stats.append(run_clip(model, row, args.videos_root, args.out_root,
                                      device, tuple(args.near_percentiles),
                                      args.resolution_level, args.fps,
                                      args.skip_existing, args.save_npz))
        except Exception as exc:
            print(f"  FAILED {row['uuid']}: {type(exc).__name__}: {exc}")
            all_stats.append({"uuid": row["uuid"],
                              "error": f"{type(exc).__name__}: {exc}"})

    args.out_root.mkdir(parents=True, exist_ok=True)
    out = args.out_root / f"summary_shard{i}of{nsh}.json"
    out.write_text(json.dumps({"model": args.model,
                               "resolution_level": args.resolution_level,
                               "clips": all_stats}, indent=2))
    print(f"wrote {out}")


if __name__ == "__main__":
    main()