scripts/moge_eval_gt.py

#!/usr/bin/env python
"""Evaluate MoGe-2 per-frame depth against real PointWorld ``scene_flows`` ground truth.

Run with the ``fpgm`` conda env (this needs ``fpgm.geometry`` + ``h5py``). MoGe
inference is a *separate* prior step (``scripts/moge_infer_episode.py``, run
under the ``moge`` env) whose ``.npz`` outputs this script only reads.

GT frame-convention detection is reused verbatim from
``fpgm.geometry.convention.detect_scene_flow_convention`` -- the same detector
``scripts/export_real_depth_video.py`` uses -- rather than re-derived here.

Only pixels covered by BOTH ground truth (visible & depth-valid scene_flows,
z-buffered into the frame) AND MoGe's own predicted mask are scored. GT
coverage itself is sparse and partial (measured previously at 58/179 frames,
34.4% of pixels on those frames) -- this script recomputes and reports that
coverage number again as a sanity check that it agrees with the earlier claim.

Metrics, computed pooled over all scored pixels across all measured frames:
    - AbsRel = mean(|pred - gt| / gt)
    - RMSE (metres)
    - delta1.25 = fraction with max(pred/gt, gt/pred) < 1.25
Reported twice per frame set:
    (a) "aligned"  -- after a per-frame least-squares scale+shift fit
                      (gt ~= a * pred + b) solved on that frame's scored pixels.
    (b) "raw"      -- no alignment, using MoGe's metric-scale output directly.
Temporal consistency: std of the per-frame fitted scale `a` across measured
frames (aligned case only) -- large std means a flickery, not video-consistent,
depth model.

Usage:
    PYTHONPATH=src /home/quang/miniconda3/envs/fpgm/bin/python scripts/moge_eval_gt.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 \\
        --out-dir outputs/moge
"""

from __future__ import annotations

import argparse
import io
import json
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:
    """Identical helper to ``export_real_depth_video.py`` -- generic h5 parsing, not GT logic."""
    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:
    """Nearest-wins splat of projected points into an (H, W) float32 GT depth map (0 = unmeasured)."""
    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)  # far first, so nearer points overwrite -> nearest-wins
    out[v[order], u[order]] = z[order]
    return out


def _fit_scale_shift(pred: np.ndarray, gt: np.ndarray) -> tuple[float, float]:
    """Closed-form least-squares (a, b) minimizing sum((a*pred + b - gt)^2)."""
    A = np.stack([pred, np.ones_like(pred)], axis=1)
    sol, *_ = np.linalg.lstsq(A, gt, rcond=None)
    return float(sol[0]), float(sol[1])


def _depth_metrics(pred: np.ndarray, gt: np.ndarray) -> dict:
    absrel = float(np.mean(np.abs(pred - gt) / gt))
    rmse = float(np.sqrt(np.mean((pred - gt) ** 2)))
    ratio = np.maximum(pred / gt, gt / pred)
    delta125 = float(np.mean(ratio < 1.25))
    return {"AbsRel": absrel, "RMSE_m": rmse, "delta1.25": delta125, "n_pixels": int(pred.size)}


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("--out-dir", required=True, type=Path)
    ap.add_argument("--min-pixels-per-frame", type=int, default=50)
    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))
    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
        print(f"frame convention: {det.convention.name} (margin={det.margin:.4f}, "
              f"frac_in_bounds={det.frac_in_bounds:.4f})")

        # GT depth is z-buffered at MoGe's native video resolution so pixels line up 1:1
        # with MoGe's per-frame depth maps (no resizing/interpolation on either side).
        gt_depth = np.zeros((n_frames, vid_h, vid_w), dtype=np.float32)
        measured = np.zeros(n_frames, dtype=bool)
        for clip in clips:
            a, b = (int(x) for x in clip.split(":"))
            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 i, t in enumerate(range(a, min(b, n_frames))):
                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)
                measured[t] = True

    measured_frames = [t for t in range(n_frames) if measured[t]]
    print(f"GT measured frames: {len(measured_frames)}/{n_frames} = {len(measured_frames) / n_frames:.1%}")

    per_frame_records = []
    pooled_pred_aligned, pooled_gt_aligned = [], []
    pooled_pred_raw, pooled_gt_raw = [], []

    for t in measured_frames:
        npz_path = args.moge_dir / f"frame_{t:03d}.npz"
        if not npz_path.exists():
            print(f"  WARNING: missing MoGe output for frame {t}, skipping")
            continue
        d = np.load(npz_path)
        moge_depth = d["depth"]
        moge_mask = d["mask"]

        gt_valid = gt_depth[t] > 0
        moge_valid = moge_mask & np.isfinite(moge_depth) & (moge_depth > 0)
        both = gt_valid & moge_valid
        n_both = int(both.sum())
        if n_both < args.min_pixels_per_frame:
            print(f"  frame {t}: only {n_both} overlapping pixels, skipping (< {args.min_pixels_per_frame})")
            continue

        pred = moge_depth[both].astype(np.float64)
        gt = gt_depth[t][both].astype(np.float64)

        a, b = _fit_scale_shift(pred, gt)
        pred_aligned = a * pred + b
        keep = pred_aligned > 0  # guard against a pathological negative-depth fit
        m_aligned = _depth_metrics(pred_aligned[keep], gt[keep])
        m_raw = _depth_metrics(pred, gt)

        per_frame_records.append({
            "frame": t,
            "n_gt_pixels": int(gt_valid.sum()),
            "n_overlap_pixels": n_both,
            "pixel_coverage_frac_gt_frame": float(gt_valid.mean()),
            "fit_scale_a": a,
            "fit_shift_b": b,
            "aligned": m_aligned,
            "raw": m_raw,
        })
        pooled_pred_aligned.append(pred_aligned[keep])
        pooled_gt_aligned.append(gt[keep])
        pooled_pred_raw.append(pred)
        pooled_gt_raw.append(gt)

    pooled_pred_aligned = np.concatenate(pooled_pred_aligned)
    pooled_gt_aligned = np.concatenate(pooled_gt_aligned)
    pooled_pred_raw = np.concatenate(pooled_pred_raw)
    pooled_gt_raw = np.concatenate(pooled_gt_raw)

    pooled_aligned_metrics = _depth_metrics(pooled_pred_aligned, pooled_gt_aligned)
    pooled_raw_metrics = _depth_metrics(pooled_pred_raw, pooled_gt_raw)

    scales = np.array([r["fit_scale_a"] for r in per_frame_records])
    temporal = {
        "n_frames_scored": len(per_frame_records),
        "fit_scale_mean": float(scales.mean()),
        "fit_scale_std": float(scales.std()),
        "fit_scale_cv": float(scales.std() / scales.mean()) if scales.mean() != 0 else float("nan"),
        "fit_scale_min": float(scales.min()),
        "fit_scale_max": float(scales.max()),
    }

    pixel_cov = np.mean([r["pixel_coverage_frac_gt_frame"] for r in per_frame_records])

    summary = {
        "video": str(args.video),
        "flows": str(args.flows),
        "camera_serial": args.camera_serial,
        "convention": det.convention.name,
        "n_frames_total": n_frames,
        "n_frames_gt_measured": len(measured_frames),
        "frame_coverage_frac": len(measured_frames) / n_frames,
        "mean_pixel_coverage_on_measured_frames": float(pixel_cov),
        "n_frames_scored_moge_overlap": len(per_frame_records),
        "pooled_aligned": pooled_aligned_metrics,
        "pooled_raw": pooled_raw_metrics,
        "temporal_scale_consistency": temporal,
        "per_frame": per_frame_records,
    }

    args.out_dir.mkdir(parents=True, exist_ok=True)
    (args.out_dir / "eval_metrics.json").write_text(json.dumps(summary, indent=2))

    print("\n=== POOLED METRICS (all scored pixels, all measured frames) ===")
    print(f"frames scored: {len(per_frame_records)}/{len(measured_frames)} GT-measured frames "
          f"(rest had <{args.min_pixels_per_frame} overlap pixels with MoGe's mask)")
    print(f"mean pixel coverage on measured frames (GT only): {pixel_cov:.1%}")
    print(f"[aligned, per-frame LS scale+shift] AbsRel={pooled_aligned_metrics['AbsRel']:.4f}  "
          f"RMSE={pooled_aligned_metrics['RMSE_m']:.4f} m  delta1.25={pooled_aligned_metrics['delta1.25']:.4f}")
    print(f"[raw, metric, no alignment]         AbsRel={pooled_raw_metrics['AbsRel']:.4f}  "
          f"RMSE={pooled_raw_metrics['RMSE_m']:.4f} m  delta1.25={pooled_raw_metrics['delta1.25']:.4f}")
    print(f"temporal scale consistency: mean(a)={temporal['fit_scale_mean']:.4f}  "
          f"std(a)={temporal['fit_scale_std']:.4f}  CV={temporal['fit_scale_cv']:.4f}  "
          f"range=[{temporal['fit_scale_min']:.4f}, {temporal['fit_scale_max']:.4f}]")
    print(f"\nwrote {args.out_dir}/eval_metrics.json")


if __name__ == "__main__":
    main()