scripts/moge_contact_sheet.py

#!/usr/bin/env python
"""Build a contact sheet: RGB / MoGe depth / GT measured points / error map, for a
handful of frames spread across the episode.

Run under ``fpgm`` env. Reuses the same GT projection + convention-detection
path as ``scripts/moge_eval_gt.py`` / ``scripts/export_real_depth_video.py``;
only recomputes GT depth for the requested frames (cheap) rather than caching
the full per-episode GT volume to disk.

Usage:
    PYTHONPATH=src /home/quang/miniconda3/envs/fpgm/bin/python scripts/moge_contact_sheet.py \\
        --flows data/pointworld/droid/flows-fs-optimized/<uuid>_flows.h5 \\
        --camera-serial 22008760 \\
        --moge-dir outputs/moge/per_frame \\
        --video data/droid_raw/<uuid>/recordings/MP4/22008760.mp4 \\
        --frames 5 25 45 80 \\
        --out outputs/moge/contact_sheet.png
"""

from __future__ import annotations

import argparse
import io
import re
from pathlib import Path

import cv2
import h5py
import numpy as np
from PIL import Image

from fpgm.config import ConventionConfig
from fpgm.geometry.camera import Camera, CameraIntrinsics
from fpgm.geometry.convention import detect_scene_flow_convention
from fpgm.types import FrameConvention

_CLIP_RE = re.compile(r"^(\d+):(\d+)$")


def _decode_initial_rgb(raw: object) -> np.ndarray:
    data = raw[0] if isinstance(raw, np.ndarray) and raw.dtype == object else raw
    if isinstance(data, bytes | bytearray | np.bytes_):
        return np.asarray(Image.open(io.BytesIO(bytes(data))).convert("RGB"))
    arr = np.asarray(data)
    if arr.ndim == 1:
        return np.asarray(Image.open(io.BytesIO(arr.tobytes())).convert("RGB"))
    return arr


def _camera_from(intrinsic: np.ndarray, extrinsic: np.ndarray, w: int, h: int) -> Camera:
    K = CameraIntrinsics(
        fx=float(intrinsic[0, 0]), fy=float(intrinsic[1, 1]),
        cx=float(intrinsic[0, 2]), cy=float(intrinsic[1, 2]), width=w, height=h,
    )
    return Camera(K, np.asarray(extrinsic, dtype=np.float64))


def _zbuffer_depth(uv: np.ndarray, z: np.ndarray, h: int, w: int) -> np.ndarray:
    out = np.zeros((h, w), dtype=np.float32)
    u = np.round(uv[:, 0]).astype(np.int64)
    v = np.round(uv[:, 1]).astype(np.int64)
    keep = (u >= 0) & (u < w) & (v >= 0) & (v < h) & (z > 0)
    u, v, z = u[keep], v[keep], z[keep]
    order = np.argsort(-z)
    out[v[order], u[order]] = z[order]
    return out


def _colorize(depth: np.ndarray, valid: np.ndarray, lo: float, hi: float, dilate: int = 0) -> np.ndarray:
    """``dilate>0`` thickens sparse single-pixel z-buffer splats for visibility only
    (metrics are always computed on the raw, non-dilated pixels elsewhere)."""
    g = np.zeros(depth.shape, dtype=np.float32)
    g[valid] = np.clip((depth[valid] - lo) / max(hi - lo, 1e-9), 0.0, 1.0)
    col = cv2.applyColorMap((g * 255).astype(np.uint8), cv2.COLORMAP_TURBO)
    col[~valid] = (30, 30, 30)
    if dilate > 0:
        kernel = np.ones((dilate, dilate), np.uint8)
        col = cv2.dilate(col, kernel)
        valid_d = cv2.dilate(valid.astype(np.uint8), kernel).astype(bool)
        col[~valid_d] = (30, 30, 30)
    return col


def _label(img: np.ndarray, text: str) -> np.ndarray:
    img = img.copy()
    cv2.rectangle(img, (0, 0), (min(len(text) * 11 + 10, img.shape[1]), 26), (0, 0, 0), -1)
    cv2.putText(img, text, (5, 19), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA)
    return img


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--flows", required=True, type=Path)
    ap.add_argument("--camera-serial", required=True)
    ap.add_argument("--moge-dir", required=True, type=Path)
    ap.add_argument("--video", required=True, type=Path)
    ap.add_argument("--frames", type=int, nargs="+", required=True)
    ap.add_argument("--out", required=True, type=Path)
    args = ap.parse_args()

    cap = cv2.VideoCapture(str(args.video))
    vid_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    vid_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    rgb_by_frame = {}
    t_cur = 0
    want = set(args.frames)
    while t_cur <= max(args.frames):
        ok, frame_bgr = cap.read()
        if not ok:
            break
        if t_cur in want:
            rgb_by_frame[t_cur] = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        t_cur += 1
    cap.release()

    with h5py.File(args.flows, "r") as f:
        clips = sorted((k for k in f if _CLIP_RE.fullmatch(k)), key=lambda s: int(s.split(":")[0]))
        n_frames = int(f.attrs.get("trajectory_length", 0)) or max(int(c.split(":")[1]) for c in clips)
        cam_key = f"camera_{args.camera_serial}_ext"

        first = f[clips[0]][cam_key]
        ann_h, ann_w = np.asarray(first["initial_depth"]).shape
        rgb0 = _decode_initial_rgb(first["initial_rgb"][()])
        cam_native = _camera_from(first["intrinsic"], first["extrinsic"], ann_w, ann_h)
        det = detect_scene_flow_convention(
            cam_native.rescaled(rgb0.shape[1], rgb0.shape[0]),
            np.asarray(first["scene_flows"][0], dtype=np.float64),
            np.asarray(first["scene_colors"][0]),
            rgb0,
            ConventionConfig(),
        )
        is_world = det.convention is FrameConvention.WORLD

        gt_depth = {t: np.zeros((vid_h, vid_w), dtype=np.float32) for t in args.frames}
        for clip in clips:
            a, b = (int(x) for x in clip.split(":"))
            relevant = [t for t in args.frames if a <= t < min(b, n_frames)]
            if not relevant:
                continue
            g = f[clip][cam_key]
            cam = _camera_from(g["intrinsic"], g["extrinsic"], ann_w, ann_h).rescaled(vid_w, vid_h)
            pts = np.asarray(g["scene_flows"], dtype=np.float64)
            vis = np.asarray(g["scene_visibility"])
            ok = np.asarray(g["scene_depth_valid_mask"])
            for t in relevant:
                i = t - a
                sel = vis[i] & ok[i]
                if not sel.any():
                    continue
                p = pts[i][sel]
                uv, z = (cam.project(p) if is_world else cam.project_cam(p))
                gt_depth[t] = _zbuffer_depth(np.asarray(uv), np.asarray(z), vid_h, vid_w)

    rows = []
    for t in args.frames:
        if t not in rgb_by_frame:
            print(f"WARNING: frame {t} not decodable from video, skipping")
            continue
        npz_path = args.moge_dir / f"frame_{t:03d}.npz"
        d = np.load(npz_path)
        moge_depth = d["depth"]
        moge_mask = d["mask"] & np.isfinite(moge_depth) & (moge_depth > 0)

        gt_valid = gt_depth[t] > 0
        both = gt_valid & moge_mask

        # Per-frame least-squares scale+shift fit (same as moge_eval_gt.py) for a fair
        # visual comparison -- otherwise MoGe's own arbitrary metric offset dominates the colormap.
        pred = moge_depth[both].astype(np.float64)
        gt = gt_depth[t][both].astype(np.float64)
        if pred.size >= 10:
            A = np.stack([pred, np.ones_like(pred)], axis=1)
            (a_fit, b_fit), *_ = np.linalg.lstsq(A, gt, rcond=None)
        else:
            a_fit, b_fit = 1.0, 0.0
        moge_depth_aligned = a_fit * moge_depth + b_fit

        lo, hi = np.percentile(gt_depth[t][gt_valid], (1, 99)) if gt_valid.any() else (0, 1)

        rgb_panel = cv2.cvtColor(rgb_by_frame[t], cv2.COLOR_RGB2BGR)
        moge_panel = _colorize(moge_depth_aligned, moge_mask, lo, hi)
        gt_panel = _colorize(gt_depth[t], gt_valid, lo, hi, dilate=5)

        err = np.zeros(gt_depth[t].shape, dtype=np.float32)
        err[both] = np.abs(moge_depth_aligned[both] - gt_depth[t][both])
        err_hi = np.percentile(err[both], 95) if both.any() else 1.0
        err_panel = _colorize(err, both, 0.0, err_hi, dilate=5)

        n_overlap = int(both.sum())
        row = np.concatenate([
            _label(rgb_panel, f"f{t:03d} RGB"),
            _label(moge_panel, f"MoGe depth (aligned a={a_fit:.2f})"),
            _label(gt_panel, "GT measured points"),
            _label(err_panel, f"|err| n={n_overlap} hi={err_hi:.3f}m"),
        ], axis=1)
        rows.append(row)

    sheet = np.concatenate(rows, axis=0)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(args.out), sheet)
    print(f"wrote {args.out}  shape={sheet.shape}")


if __name__ == "__main__":
    main()