scripts/moge_infer_episode.py

#!/usr/bin/env python
"""Run Microsoft MoGe-2 on every frame of a DROID episode mp4.

Standalone script -- must be run with the ``moge`` conda env
(``/home/quang/miniconda3/envs/moge/bin/python``), NOT ``fpgm``. It has no
dependency on the ``fpgm`` package, only on ``moge`` (third_party/MoGe) and
plain ``torch``/``cv2``/``numpy``.

For each decodable frame it saves a compressed ``.npz`` with:
    depth: (H, W) float32, metres, camera-frame Z (forward), MoGe's native
           output convention -- directly comparable to this project's GT
           ``Camera.project()`` depth (also camera-frame Z).
    mask:  (H, W) bool, MoGe's own predicted valid-pixel mask.
    intrinsics: (3, 3) float32, MoGe's *own estimated* normalized intrinsics
           (fov not given -- this is the out-of-the-box usage path).

Usage:
    /home/quang/miniconda3/envs/moge/bin/python scripts/moge_infer_episode.py \\
        --video data/droid_raw/<uuid>/recordings/MP4/22008760.mp4 \\
        --out-dir outputs/moge/per_frame \\
        --model Ruicheng/moge-2-vitl-normal
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import cv2
import numpy as np
import torch


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--video", required=True, type=Path)
    ap.add_argument("--out-dir", required=True, type=Path)
    ap.add_argument("--model", default="Ruicheng/moge-2-vitl-normal")
    ap.add_argument("--device", default="cuda")
    args = ap.parse_args()

    from moge.model.v2 import MoGeModel

    device = torch.device(args.device)
    print(f"loading {args.model} ...")
    model = MoGeModel.from_pretrained(args.model).to(device).eval()

    cap = cv2.VideoCapture(str(args.video))
    n_frames = 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))
    print(f"video: {args.video} -- {n_frames} frames, {w}x{h}")

    args.out_dir.mkdir(parents=True, exist_ok=True)

    n_done = 0
    t0 = time.time()
    with torch.inference_mode():
        for t in range(n_frames):
            ok, frame_bgr = cap.read()
            if not ok:
                print(f"WARNING: video reported {n_frames} frames but decoding stopped at t={t}")
                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)

            output = model.infer(image)  # out-of-the-box: no fov_x given, MoGe estimates its own FOV
            depth = output["depth"].detach().float().cpu().numpy()
            mask = output["mask"].detach().cpu().numpy().astype(bool)
            intrinsics = output["intrinsics"].detach().float().cpu().numpy()

            # infinite depth outside MoGe's own predicted mask -> store as NaN, not inf,
            # so downstream numpy nan-handling is explicit rather than accidental.
            depth = np.where(np.isfinite(depth), depth, np.nan).astype(np.float32)

            np.savez_compressed(
                args.out_dir / f"frame_{t:03d}.npz",
                depth=depth,
                mask=mask,
                intrinsics=intrinsics,
            )
            n_done += 1
            if n_done % 20 == 0 or n_done == 1:
                print(f"  {n_done}/{n_frames} done ({time.time() - t0:.1f}s elapsed)")

    cap.release()
    meta = {
        "video": str(args.video),
        "model": args.model,
        "n_frames_reported_by_container": n_frames,
        "n_frames_processed": n_done,
        "width": w,
        "height": h,
        "elapsed_sec": time.time() - t0,
        "fov_x_given": False,
        "note": "depth is MoGe's own camera-frame Z (points[...,2]); metric_scale applied since model is a metric checkpoint.",
    }
    (args.out_dir / "run_meta.json").write_text(json.dumps(meta, indent=2))
    print(f"done: {n_done} frames in {meta['elapsed_sec']:.1f}s -> {args.out_dir}")


if __name__ == "__main__":
    main()