scripts/export_depth_gt_video.py

#!/usr/bin/env python
"""Whole-episode ground-truth depth as viewable video: inverse-depth + provenance.

Two things this does that ``master/depth_preview.mp4`` does not, both because it
reads ``master/depth_dense.h5`` (the metric ``uint16`` millimetres) rather than a
pre-rendered preview:

**One global normalisation across the whole episode, never per frame.** Per-frame
normalisation makes a depth video flicker as the scene's near/far extremes move,
and this project has already recorded the consequence downstream: a model trained
on it reproduces the flicker as brightness pumping (see
``scripts/export_cosmos_input.py``'s "Depth is normalised once for the whole clip"
note and ``export_vace``'s ``_INVERSE_DEPTH_PERCENTILES``). It also matters for
*looking* at the data: with per-frame ranges, the eye cannot tell a surface that
moved from a normalisation that shifted. Robust percentiles, not min/max, so a
single stray pixel cannot crush the usable range.

**A provenance track.** ``depth_dense.h5``'s ``source`` array records, per pixel
per frame, why that depth has the value it has (:class:`DepthSourceCode`):
``SPLAT`` = a real projected scene_flows point *this frame*, ``ANCHOR`` = the
clip's dense initial-depth anchor, ``FILLED`` = push-pull interpolated from
neighbouring support, ``INVALID`` = no support in range and ``depth_mm`` is 0.
A depth video alone cannot show the difference between measured and interpolated
geometry -- they are both just grey -- so eyeballing depth quality without this
track systematically overrates it.

Usage:
    PYTHONPATH=src python scripts/export_depth_gt_video.py \\
        --master outputs/datagen/<uuid>/<cam>/master \\
        --out-dir outputs/depth_gt [--scale 3]
"""

from __future__ import annotations

import argparse
import subprocess
from pathlib import Path

import cv2
import h5py
import numpy as np

from fpgm.datagen.types import DepthSourceCode
from fpgm.viz.video import _resolve_ffmpeg

#: Robust range for the single global inverse-depth normalisation.
_PERCENTILES = (1.0, 99.0)

#: Provenance palette, BGR. Chosen so the two that mean "measured" (SPLAT,
#: ANCHOR) read as one family and the two that mean "not measured" (FILLED,
#: INVALID) as another -- the distinction that actually matters when judging
#: whether a region of the depth map is evidence or inference.
_SOURCE_BGR = {
    DepthSourceCode.SPLAT: (80, 220, 80),      # green  -- measured this frame
    DepthSourceCode.ANCHOR: (220, 200, 60),    # cyan   -- measured at the anchor
    DepthSourceCode.FILLED: (60, 140, 240),    # orange -- interpolated
    DepthSourceCode.INVALID: (40, 40, 220),    # red    -- nothing at all
}


def write_h264(frames_bgr: np.ndarray, out: Path, fps: int) -> None:
    t, h, w, _ = frames_bgr.shape
    out.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        _resolve_ffmpeg(), "-y", "-hide_banner", "-loglevel", "error",
        "-f", "rawvideo", "-pix_fmt", "bgr24", "-s", f"{w}x{h}", "-r", str(fps),
        "-i", "pipe:0", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "16",
        str(out),
    ]
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE)
    assert p.stdin is not None
    try:
        for i in range(t):
            p.stdin.write(np.ascontiguousarray(frames_bgr[i]).tobytes())
    finally:
        p.stdin.close()
        if p.wait() != 0:
            raise RuntimeError(f"ffmpeg failed writing {out}")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--master", required=True, type=Path)
    ap.add_argument("--out-dir", required=True, type=Path)
    ap.add_argument("--scale", type=int, default=3, help="nearest-neighbour upscale for viewing")
    ap.add_argument("--fps", type=int, default=16)
    args = ap.parse_args()

    with h5py.File(args.master / "depth_dense.h5", "r") as f:
        depth_mm = np.asarray(f["depth_mm"])          # (T, H, W) uint16
        source = np.asarray(f["source"])              # (T, H, W) uint8
        support_frac = np.asarray(f["support_frac"])  # (T,)

    valid = depth_mm > 0
    depth_m = depth_mm.astype(np.float32) / 1000.0
    inv = np.zeros_like(depth_m)
    inv[valid] = 1.0 / np.maximum(depth_m[valid], 1e-6)
    lo, hi = np.percentile(inv[valid], _PERCENTILES)
    if not hi > lo:
        hi = lo + 1e-6
    gray = np.clip((inv - lo) / (hi - lo), 0.0, 1.0)
    gray[~valid] = 0.0
    gray8 = (gray * 255.0).astype(np.uint8)

    t, h, w = gray8.shape
    s = max(1, args.scale)
    out_h, out_w = h * s, w * s

    depth_frames = np.empty((t, out_h, out_w, 3), dtype=np.uint8)
    source_frames = np.empty((t, out_h, out_w, 3), dtype=np.uint8)
    for i in range(t):
        up = cv2.resize(gray8[i], (out_w, out_h), interpolation=cv2.INTER_NEAREST)
        depth_frames[i] = cv2.applyColorMap(up, cv2.COLORMAP_TURBO)
        rgb = np.zeros((h, w, 3), dtype=np.uint8)
        for code, bgr in _SOURCE_BGR.items():
            rgb[source[i] == int(code)] = bgr
        source_frames[i] = cv2.resize(rgb, (out_w, out_h), interpolation=cv2.INTER_NEAREST)

    args.out_dir.mkdir(parents=True, exist_ok=True)
    write_h264(depth_frames, args.out_dir / "depth_gt_full.mp4", args.fps)
    write_h264(source_frames, args.out_dir / "depth_source_full.mp4", args.fps)
    write_h264(
        np.concatenate([depth_frames, source_frames], axis=2),
        args.out_dir / "depth_gt_and_source.mp4", args.fps,
    )

    counts = {c.name: float((source == int(c)).mean()) for c in DepthSourceCode}
    print(f"frames:            {t}  ({h}x{w} native, upscaled x{s})")
    print(f"valid depth px:    {valid.mean():.2%}")
    dmin = depth_m[valid].min() if valid.any() else float("nan")
    dmax = depth_m[valid].max() if valid.any() else float("nan")
    print(f"metric range:      {dmin:.3f} .. {dmax:.3f} m")
    print(f"global inv range:  {lo:.4f} .. {hi:.4f} per-m (p{_PERCENTILES[0]}..p{_PERCENTILES[1]})")
    print("provenance (fraction of all pixels):")
    for name, frac in counts.items():
        print(f"  {name:<8} {frac:7.2%}")
    print(f"support_frac (non-FILLED) per frame: min {support_frac.min():.3f} "
          f"mean {support_frac.mean():.3f} max {support_frac.max():.3f}")
    print(f"wrote: {args.out_dir}/depth_gt_full.mp4, depth_source_full.mp4, depth_gt_and_source.mp4")


if __name__ == "__main__":
    main()