scripts/publish_moge_depth.py

#!/usr/bin/env python
"""Publish the MoGe-2 depth preview videos to a subfolder of the datagen bucket.

The dataset itself is published by ``scripts/publish_dataset.py`` /
:class:`fpgm.datagen.publish.DatasetPublisher`, which syncs the staged episode
tiers to the ROOT of ``PublishConfig.bucket_url``. This is deliberately a
separate, much smaller shipment: qualitative ``RGB | depth | normal`` previews
that live under ``<bucket>/moge_depth/`` so they cannot collide with, or be
mistaken for, the episode payload.

Token policy is copied verbatim from ``DatasetPublisher.publish`` and must stay
that way: ``HF_TOKEN`` is read from the environment and passed explicitly, and a
cached token at ``~/.cache/huggingface/token`` is NEVER consulted -- ``token=False``
is passed when ``HF_TOKEN`` is unset, which disables that fallback rather than
silently inheriting whatever identity happens to be logged in on the box.

``delete`` is not exposed at all. A sync that can delete has no business being one
flag away from a preview upload into a bucket that holds the real dataset.

Usage:
    PYTHONPATH=src python scripts/publish_moge_depth.py                  # dry-run
    HF_TOKEN=hf_... PYTHONPATH=src python scripts/publish_moge_depth.py --no-dry-run
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

SUBFOLDER = "moge_depth"


def build_staging(clips_root: Path, staging: Path, step_src: Path | None,
                  scene_note: str) -> list[dict]:
    """Copy just the preview mp4s + their stats into a flat, small staging dir.

    The per-frame ``.npz`` volumes are ~1 GB per clip and are NOT shipped: they are
    a local recompute cache, not something anyone judging the depth needs.
    """
    if staging.exists():
        shutil.rmtree(staging)
    staging.mkdir(parents=True)

    entries = []
    for d in sorted(p for p in clips_root.iterdir() if p.is_dir()):
        # panels_2x2.mp4 (moge_render_panels.py), not the 3-panel rgb_depth_normal.mp4:
        # the 2x2 carries the depth high-pass panel, which is the only one in which a
        # 17 mm object step is visible at this scene's ~1.1 m global depth range.
        vid = d / "panels_2x2.mp4"
        stats_p = d / "stats.json"
        if not vid.exists():
            print(f"  skip {d.name}: no panels_2x2.mp4 (run moge_render_panels.py first)")
            continue
        shutil.copy2(vid, staging / f"{d.name}.mp4")
        stats = json.loads(stats_p.read_text()) if stats_p.exists() else {}
        stats["preview_video"] = f"{d.name}.mp4"
        stats["preview_bytes"] = (staging / f"{d.name}.mp4").stat().st_size
        entries.append(stats)

    # The quantitative companion to the videos (scripts/moge_object_step.py). Passed
    # explicitly rather than guessed from clips_root's parent: both preview sets live
    # under outputs/moge/, so a guess would have shipped the datagen bucket's
    # object_step.json alongside the droid-sample previews it does not describe.
    step_section = ""
    if step_src is not None and step_src.exists():
        shutil.copy2(step_src, staging / "object_step.json")
        step_section = STEP_SECTION
    elif step_src is not None:
        print(f"  note: --object-step {step_src} does not exist, not shipped")

    summary_src = next((p for p in (clips_root / "summary.json",
                                    *sorted(clips_root.glob("summary_shard*.json")))
                        if p.exists()), None)
    summary = json.loads(summary_src.read_text()) if summary_src else {}
    (staging / "summary.json").write_text(json.dumps(
        {"model": summary.get("model"),
         "resolution_level": summary.get("resolution_level"),
         "clips": entries}, indent=2))
    (staging / "README.md").write_text(
        render_readme(entries, summary, scene_note, step_section))
    return entries


DATAGEN_SCENE_NOTE = """All of these are the **same scene** (AUTOLab+0d4edc83, "Put brick in drawer shelf and
close drawer"), recorded seconds apart, because that is the whole of what this bucket
holds. Ten clips agreeing therefore is **not** evidence that MoGe-2 is stable — it is
one condition measured ten times. The only genuinely different viewpoint in the set is
camera `24400334` (ext2)."""

SAMPLE_SCENE_NOTE = """These clips come from the diversity sample: different labs, different object classes,
different scene depths. That is what makes them worth looking at — and also why the
per-clip colour ranges differ so much. Do not compare colours between two clips; only
the near-field range printed in each panel makes a panel's colours mean anything."""


def render_readme(entries: list[dict], summary: dict, scene_note: str,
                  step_section: str = "") -> str:
    def rng(e: dict) -> str:
        # datagen previews carry depth_p2_m/depth_p98_m; droid-sample previews carry
        # depth_range_m. Handle both rather than printing nan for one of them.
        if "depth_range_m" in e:
            lo, hi = e["depth_range_m"]
        else:
            lo, hi = e.get("depth_p2_m", float("nan")), e.get("depth_p98_m", float("nan"))
        return f"{lo:.2f}–{hi:.2f}"

    rows = "\n".join(
        f"| `{e.get('preview_video')}` | "
        f"{e.get('camera_role') or e.get('lab') or '?'}"
        f"{' / ' + e['object_class'] if e.get('object_class') else ''} | "
        f"{e.get('n_frames_processed', '?')} | {rng(e)} | "
        f"{e.get('near_field_mm_per_colour_level', float('nan')):.1f} | "
        f"{e.get('preview_bytes', 0) / 1e6:.1f} MB |"
        for e in entries if "error" not in e)
    return f"""# MoGe-2 monocular depth — qualitative previews

Full-length monocular depth over the clips published in this bucket, rendered as a
2x2 panel video (2560x1440, 15 fps): **RGB | depth (global range) | the same depth at a
near-field range | normal**, each panel at the clip's native 1280x720.

These are previews for judging detail by eye. They are **not** a dataset tier and
nothing in the pipeline consumes them — the pipeline's own depth comes from
PointWorld `scene_flows` (sparse metric 3D tracks), not from MoGe.

## Model

`{summary.get('model')}`, `resolution_level={summary.get('resolution_level')}`.

`resolution_level=9` is the default and already maps to the **top** of MoGe-2's
`num_tokens_range` (3600 tokens) — the finest setting the model offers.

**MoGe-3 is not used, because it does not exist as runnable code.** The paper
(arXiv 2607.17967, "Fine-Detail Monocular Geometry Estimation…") was announced in
the upstream README on 2026-07-21, but at pinned commit `925b8ed` — still
`origin/main` as of this run — there is no `moge/model/v3.py` and the pretrained
table lists `moge-3-vitg`/`moge-3-vitl` as "coming soon" with no HF repo id.

## How to read the panels

* **Depth** uses **one global normalisation for the whole clip** (robust 2/98
  percentiles over every valid pixel of every frame), never per frame. Per-frame
  normalisation makes the panel breathe as the near/far extremes move, which by eye
  is indistinguishable from the model's own temporal instability — the exact thing
  these videos exist to let someone judge. The colour range in metres is burned into
  the panel label.
* **Normal** is where fine detail actually shows. Depth is dominated by the
  metre-scale front-to-back ramp of the workbench; the normal map spends its whole
  colour budget on surface relief instead.
* **No `fov_x` was passed**, so MoGe estimates its own FOV per frame even though this
  project knows the true DROID intrinsics. That is the out-of-the-box path — priming
  the model with the calibration would measure a best case that arbitrary video would
  not get.
* Videos are muxed at **15 fps**, the DROID trajectory rate. The source mp4
  containers advertise 60/1, but their frames are the 15 Hz trajectory samples (this
  repo indexes poses per video frame at `DEFAULT_TRAJECTORY_FPS = 15.0`). Playing at
  60 would look like temporal jitter that is not there.

## Clips

| file | camera / class | frames | depth range (m) | near mm/level | size |
| --- | --- | --- | --- | --- | --- |
{rows}

The **near mm/level** column is the one to read before trusting a panel: it is the
millimetres of depth per colour level in the near-field panel. Scenes with a deep
background push the *global* panel to 15–19 mm/level, where nothing object-sized is
visible at all; the near-field panel is what stays readable.

{scene_note}
{step_section}"""


STEP_SECTION = """
## `object_step.json` — the quantitative companion

Whether the model separates the manipulated object from the surface under it, measured
with the pipeline's own object mask (`master/prompt_0_masks.h5`):

* `step_mm` — mean depth of a 25 px ring around the object minus mean depth inside it.
  Positive = object in front of its surroundings.
* `edge_ratio` — median |∇depth| on the mask boundary over the median in the ring.
  1.0 = boundary indistinguishable from flat surface. This is the half of the question
  a step size cannot answer: a model that smears a real 20 mm step over 30 px still
  reports 20 mm.

Both are computed **within a single frame** on purpose. With no `fov_x` given, MoGe
estimates its own scale per frame and that scale drifts, so absolute metres are not
comparable across frames; an inside-minus-outside difference cancels a per-frame scale
factor to first order and is.

Two caveats that travel with these numbers:

* `step_mm_p10` goes negative on several clips. That is the measurement's limit, not the
  model's error: once the gripper closes on the object, the 25 px ring contains the
  fingers, which are nearer the camera than the object, so the difference flips sign.
  The median is robust to this; the p10 is not.
* Two clips have no `prompt_0_masks.h5` at all (`...07m-04s/24400334`,
  `...09m-53s/22008760`) — both are the entries with `n_win=0` in the dataset manifest,
  i.e. datagen never produced object masks for them. They ship a video and no step.
"""


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("--staging-dir", type=Path,
                    default=REPO_ROOT / "outputs/moge/_publish_moge_depth")
    ap.add_argument("--profile", default=str(REPO_ROOT / "configs/datagen_droid.yaml"))
    ap.add_argument("--bucket-url", default=None,
                    help="override the destination bucket (default: PublishConfig.bucket_url, "
                         "i.e. the datagen bucket). Use the bucket the clips came FROM.")
    ap.add_argument("--object-step", type=Path, default=None,
                    help="path to object_step.json to ship alongside; omitted if unset")
    ap.add_argument("--scene-note", choices=("datagen", "sample"), default="datagen",
                    help="which 'how to read this set' paragraph the README gets")
    ap.add_argument("--dry-run", action=argparse.BooleanOptionalAction, default=True,
                    help="plan only, no upload (default: on)")
    args = ap.parse_args()

    from fpgm.config_datagen import DatagenProfile
    from huggingface_hub import HfApi

    bucket = args.bucket_url or DatagenProfile.from_yaml(args.profile).publish.bucket_url
    dest = f"{bucket.rstrip('/')}/{SUBFOLDER}"
    scene_note = {"datagen": DATAGEN_SCENE_NOTE, "sample": SAMPLE_SCENE_NOTE}[args.scene_note]

    print(f"staging {args.clips_root} -> {args.staging_dir}")
    entries = build_staging(args.clips_root, args.staging_dir, args.object_step, scene_note)
    total = sum(e.get("preview_bytes", 0) for e in entries)
    print(f"  {len(entries)} preview(s), {total / 1e6:.1f} MB total")
    for e in entries:
        print(f"    {e.get('preview_video')}  {e.get('n_frames_processed')} frames  "
              f"{e.get('preview_bytes', 0) / 1e6:.1f} MB")

    token = os.environ.get("HF_TOKEN") or None
    if not args.dry_run and not token:
        raise RuntimeError(
            "publish_moge_depth: HF_TOKEN is not set in the environment. Refusing a real "
            "upload without an explicit token (this never falls back to a cached token at "
            "~/.cache/huggingface/token). Pass --dry-run, or set HF_TOKEN.")

    api = HfApi(token=(token if token else False))
    print(f"\nsync {args.staging_dir} -> {dest}  (dry_run={args.dry_run}, "
          f"token={'<set>' if token else '<none>'})")
    plan = api.sync_bucket(source=str(args.staging_dir), dest=dest,
                           dry_run=args.dry_run, delete=False, verbose=True,
                           token=(token if token else False))
    print(f"plan: {plan.summary()}")
    if args.dry_run:
        print("--dry-run: nothing uploaded. Re-run with HF_TOKEN=... --no-dry-run")


if __name__ == "__main__":
    main()