scripts/export_real_depth_video.py

#!/usr/bin/env python
"""Whole-episode depth built ONLY from real per-frame measurements -- nothing filled.

Answers "what does the depth actually look like before any interpolation?", which
neither existing artifact does: ``master/depth_dense.h5`` is 90.5% ANCHOR + 9.5%
push-pull FILLED with `support_frac` bit-identical across every frame (i.e. one
static map repeated), and the exported ``control_depth.mkv`` is a mesh z-buffer
composited over that static background plate. Both are useful; neither shows the
raw evidence.

The raw evidence is PointWorld's ``scene_flows``: ``(T, N, 3)`` metric 3D point
tracks, per clip, gated by ``scene_visibility`` and ``scene_depth_valid_mask``.
This projects them into the camera with a z-buffer and leaves every unmeasured
pixel black. **The gaps are the point of the output**, so they are never filled.

Two things this surfaces that a filled depth map hides:

* **PointWorld does not cover every frame.** Its clips are short overlapping
  windows (e.g. ``0:11``, ``5:16``, ``20:31``...) and an episode is typically
  covered only in part; the uncovered frames have no measured depth *at all*.
  Coverage is printed and written into the sidecar JSON.
* **Even a covered frame is sparse.** ~27k tracked points against 320x180 = 57,600
  pixels is under half the frame before visibility/validity gating.

Frame convention (WORLD vs CAMERA ``scene_flows``) is not assumed: it is decided by
``fpgm.geometry.convention.detect_scene_flow_convention``, the same detector the
pipeline itself uses, against the clip's own ``initial_rgb``/``scene_colors``.

Usage:
    PYTHONPATH=src python scripts/export_real_depth_video.py \\
        --flows data/pointworld/droid/flows-fs-optimized/<uuid>_flows.h5 \\
        --camera-serial 22008760 --out-dir outputs/depth_real [--scale 3]
"""

from __future__ import annotations

import argparse
import io
import json
import re
import subprocess
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
from fpgm.viz.video import _resolve_ffmpeg

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


def _decode_initial_rgb(raw: object) -> np.ndarray:
    """``initial_rgb`` -> ``(H, W, 3)`` uint8.

    The dataset stores it as a 1-element object array holding **JPEG bytes**, and
    h5py hands those back as a 1-D ``uint8`` ndarray rather than a ``bytes`` object,
    so an ``isinstance(..., bytes)`` check alone silently returns the raw byte
    string as if it were an image.
    """
    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:  # encoded image bytes
        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(uv: np.ndarray, z: np.ndarray, h: int, w: int) -> np.ndarray:
    """Nearest-wins splat of projected points into an (H, W) float32 depth map."""
    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
    np.maximum.at(out, (v[order], u[order]), 0)  # touch, keeps dtype/shape semantics
    out[v[order], u[order]] = z[order]
    return out


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("--flows", required=True, type=Path)
    ap.add_argument("--camera-serial", required=True)
    ap.add_argument("--out-dir", required=True, type=Path)
    ap.add_argument("--scale", type=int, default=3)
    ap.add_argument("--fps", type=int, default=16)
    args = ap.parse_args()

    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]),
        )
        if not clips:
            raise SystemExit(f"no clip groups in {args.flows}")
        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]
        h, w = np.asarray(first["initial_depth"]).shape
        rgb0 = _decode_initial_rgb(first["initial_rgb"][()])
        cam_native = _camera_from(first["intrinsic"], first["extrinsic"], w, 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}")

        depth = np.zeros((n_frames, h, 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"], w, 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))
                depth[t] = _zbuffer(np.asarray(uv), np.asarray(z), h, w)
                measured[t] = True

    valid = depth > 0
    if not valid.any():
        raise SystemExit("no measured depth recovered -- check --camera-serial")
    lo, hi = np.percentile(1.0 / depth[valid], (1.0, 99.0))
    if not hi > lo:
        hi = lo + 1e-6

    s = max(1, args.scale)
    out_h, out_w = h * s, w * s
    frames = np.zeros((n_frames, out_h, out_w, 3), dtype=np.uint8)
    for t in range(n_frames):
        g = np.zeros((h, w), dtype=np.float32)
        m = depth[t] > 0
        if m.any():
            g[m] = np.clip((1.0 / depth[t][m] - lo) / (hi - lo), 0.0, 1.0)
        up = cv2.resize((g * 255).astype(np.uint8), (out_w, out_h), interpolation=cv2.INTER_NEAREST)
        col = cv2.applyColorMap(up, cv2.COLORMAP_TURBO)
        col[cv2.resize(m.astype(np.uint8), (out_w, out_h), interpolation=cv2.INTER_NEAREST) == 0] = 0
        tag = f"f{t:03d}  " + ("MEASURED" if measured[t] else "NO POINTWORLD DATA")
        cv2.rectangle(col, (0, 0), (len(tag) * 11 + 10, 26), (0, 0, 0), -1)
        cv2.putText(col, tag, (5, 19), cv2.FONT_HERSHEY_SIMPLEX, 0.6,
                    (255, 255, 255) if measured[t] else (60, 60, 255), 1, cv2.LINE_AA)
        frames[t] = col

    args.out_dir.mkdir(parents=True, exist_ok=True)
    write_h264(frames, args.out_dir / "depth_real_full.mp4", args.fps)

    per_frame_cov = valid.reshape(n_frames, -1).mean(axis=1)
    stats = {
        "n_frames": int(n_frames),
        "frames_with_any_measurement": int(measured.sum()),
        "frame_coverage": float(measured.mean()),
        "pixel_coverage_on_measured_frames": float(per_frame_cov[measured].mean()) if measured.any() else 0.0,
        "clips": clips,
        "convention": det.convention.name,
        "metric_range_m": [float(depth[valid].min()), float(depth[valid].max())],
    }
    (args.out_dir / "depth_real_stats.json").write_text(json.dumps(stats, indent=2))
    print(f"frames with real measurement: {measured.sum()}/{n_frames} = {measured.mean():.1%}")
    print(f"pixel coverage on those frames: {stats['pixel_coverage_on_measured_frames']:.1%}")
    print(f"metric range: {stats['metric_range_m'][0]:.3f} .. {stats['metric_range_m'][1]:.3f} m")
    print(f"wrote {args.out_dir}/depth_real_full.mp4")


if __name__ == "__main__":
    main()