scripts/moge_object_step.py

#!/usr/bin/env python
"""How big a depth step does MoGe-2 actually put on the manipulated object?

The videos let someone judge detail by eye; this puts a number on the one thing that
matters for this project, which is whether the model separates the *manipulated
object* from the surface it sits on. Both quantities are measured **within a single
frame**, deliberately:

MoGe is run out-of-the-box with no `fov_x`, so it estimates its own FOV and its own
metric scale per frame, and that scale drifts (CV ~9.6% on this project's demo
episode, `outputs/moge/README.md`). Any absolute depth in metres is therefore not
comparable across frames. A step measured inside one frame is immune to that drift,
because a per-frame scale factor cancels out of an inside-minus-outside difference
to first order.

Per frame, using the pipeline's own object mask (`master/prompt_0_masks.h5`, the
same mask the velocity pipeline tracks):

* `step_mm`   -- mean depth of a dilated ring around the object minus mean depth
                 inside it. Positive = the object is IN FRONT of its surroundings,
                 which is what a real object on a surface must be.
* `edge_ratio`-- median |grad depth| on the mask boundary divided by the median in
                 the ring. This is the sharpness half of the question and the half
                 a step size cannot answer: a model that smears a real 17 mm step
                 across 30 px still reports 17 mm. 1.0 means the boundary is
                 indistinguishable from flat surface; higher is crisper.

Frames whose mask is empty (object out of view / occluded) are skipped, and the
count of those is reported rather than silently dropped -- an episode where the
object is visible in 40% of frames is a different measurement from one where it is
visible throughout.

Env: anything with numpy/cv2/h5py (`fpgm`).

Usage:
    /home/quang/miniconda3/envs/fpgm/bin/python scripts/moge_object_step.py \\
        [--clips-root outputs/moge/published_clips] [--out outputs/moge/object_step.json]
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import cv2
import h5py
import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
STAGING = REPO_ROOT / "outputs/datagen/_publish_staging"


def measure_clip(clip_dir: Path, ring_px: int, stride: int) -> dict:
    uuid, serial = clip_dir.name.rsplit("__", 1)
    masks_h5 = STAGING / uuid / serial / "master" / "prompt_0_masks.h5"
    if not masks_h5.exists():
        return {"clip": clip_dir.name, "error": f"no mask file at {masks_h5}"}

    files = sorted((clip_dir / "per_frame").glob("frame_*.npz"))
    with h5py.File(masks_h5, "r") as f:
        masks = f["mask"]
        n = min(len(files), masks.shape[0])
        ring_k = np.ones((ring_px, ring_px), np.uint8)
        inner_k = np.ones((5, 5), np.uint8)
        edge_k = np.ones((3, 3), np.uint8)

        steps, ratios, areas = [], [], []
        n_empty = 0
        for t in range(0, n, stride):
            m = masks[t].astype(np.uint8)
            if m.sum() == 0:
                n_empty += 1
                continue
            depth = np.load(files[t])["depth"].astype(np.float32)
            finite = np.isfinite(depth)
            ring = cv2.dilate(m, ring_k) - cv2.dilate(m, inner_k)
            inside = depth[(m > 0) & finite]
            outside = depth[(ring > 0) & finite]
            if inside.size < 50 or outside.size < 50:
                n_empty += 1
                continue
            steps.append((outside.mean() - inside.mean()) * 1000.0)
            areas.append(int(m.sum()))

            gy, gx = np.gradient(np.where(finite, depth, np.nan))
            g = np.hypot(gx, gy)
            edge = cv2.dilate(m, edge_k) - cv2.erode(m, edge_k)
            ge = g[edge > 0]
            gf = g[ring > 0]
            ge, gf = ge[np.isfinite(ge)], gf[np.isfinite(gf)]
            if ge.size and gf.size and np.median(gf) > 0:
                ratios.append(float(np.median(ge) / np.median(gf)))

    if not steps:
        return {"clip": clip_dir.name, "n_frames_measured": 0,
                "n_frames_skipped": n_empty, "note": "object never usable in mask"}
    return {
        "clip": clip_dir.name, "uuid": uuid, "camera_serial": serial,
        "n_frames_total": n, "n_frames_measured": len(steps),
        "n_frames_skipped_empty_mask": n_empty,
        "object_area_px_median": int(np.median(areas)),
        "step_mm_median": round(float(np.median(steps)), 2),
        "step_mm_p10": round(float(np.percentile(steps, 10)), 2),
        "step_mm_p90": round(float(np.percentile(steps, 90)), 2),
        "step_mm_frac_positive": round(float(np.mean(np.array(steps) > 0)), 3),
        "edge_ratio_median": round(float(np.median(ratios)), 3) if ratios else None,
    }


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("--ring-px", type=int, default=25,
                    help="dilation size of the comparison ring around the object mask")
    ap.add_argument("--stride", type=int, default=2, help="measure every Nth frame")
    ap.add_argument("--out", type=Path, default=REPO_ROOT / "outputs/moge/object_step.json")
    args = ap.parse_args()

    rows = []
    for d in sorted(p for p in args.clips_root.iterdir() if p.is_dir()):
        if not (d / "per_frame").exists():
            continue
        r = measure_clip(d, args.ring_px, args.stride)
        rows.append(r)
        if "error" in r or not r.get("n_frames_measured"):
            print(f"{d.name}: {r.get('error') or r.get('note')}")
        else:
            print(f"{r['clip'][-24:]:>24}  step {r['step_mm_median']:+7.1f} mm "
                  f"[p10 {r['step_mm_p10']:+.1f}, p90 {r['step_mm_p90']:+.1f}]  "
                  f"edge x{r['edge_ratio_median']:.2f}  "
                  f"area {r['object_area_px_median']:5d} px  "
                  f"n={r['n_frames_measured']} (+{r['n_frames_skipped_empty_mask']} skipped)")

    ok = [r for r in rows if r.get("n_frames_measured")]
    if ok:
        allsteps = [r["step_mm_median"] for r in ok]
        allratio = [r["edge_ratio_median"] for r in ok if r["edge_ratio_median"]]
        print(f"\nacross {len(ok)} clip(s): step median {np.median(allsteps):+.1f} mm "
              f"(range {min(allsteps):+.1f}..{max(allsteps):+.1f}), "
              f"edge ratio median x{np.median(allratio):.2f}")
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps({"ring_px": args.ring_px, "stride": args.stride,
                                    "clips": rows}, indent=2))
    print(f"wrote {args.out}")


if __name__ == "__main__":
    main()