scripts/export_moge_pointcloud_control.py

#!/usr/bin/env python
"""Build a VACE control video by re-rendering a clip's MoGe-2 point cloud.

Env + objects come entirely from MoGe: every pixel is unprojected with MoGe's own
per-frame estimated intrinsics, then z-buffer-splatted back into the same camera.
Colour is the real frame's RGB carried on the points, so what reaches VACE is
photograph-like rather than a false-colour geometry buffer.

That last point is the whole design constraint, and it is measured, not stylistic:
``wan/vace.py::vace_encode_frames`` runs ``vae.encode()`` on the control, so the
control passes through Wan's natural-video VAE before the DiT sees it. Feeding this
project's geometry buffers straight in scored PSNR 9.2-9.9 with the generated clip's
temporal activity pinned at the *control's* (0.0056 vs the real 0.0131) -- i.e. pure
passthrough. Compositing something photograph-like instead scored 17.0-18.0.

What the splat gaps mean: they are the honest signature of a point cloud rendered
from a single view -- disocclusions and depth-discontinuity thinning. They are NOT
filled, because filling them would hide exactly the thing this probe is meant to
show (how complete MoGe's geometry is).

Usage:
    /home/quang/miniconda3/envs/fpgm/bin/python scripts/export_moge_pointcloud_control.py \\
        --clip-dir outputs/moge/sample_bucket/<uuid> --out outputs/zeroshot_moge/<uuid>
"""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
from pathlib import Path

import cv2
import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
_FFMPEG = Path("/home/quang/miniconda3/envs/ffmpeg_libs/bin/ffmpeg")


def ffmpeg() -> str:
    return str(_FFMPEG) if _FFMPEG.exists() else (shutil.which("ffmpeg") or "ffmpeg")


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


def splat(depth: np.ndarray, rgb: np.ndarray, K: np.ndarray, radius: int) -> tuple[np.ndarray, np.ndarray]:
    """Unproject every valid pixel and z-buffer-splat it back through the same K.

    Returns ``(rendered_bgr, filled_mask)``. Nearest point wins per output pixel;
    ``radius`` dilates each point into a square so the cloud reads as a surface
    rather than a stipple, without inventing geometry between disconnected points.
    """
    h, w = depth.shape
    valid = np.isfinite(depth) & (depth > 0)
    ys, xs = np.nonzero(valid)
    z = depth[ys, xs].astype(np.float32)
    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
    X = (xs - cx) / fx * z
    Y = (ys - cy) / fy * z

    u = np.round(fx * X / z + cx).astype(np.int32)
    v = np.round(fy * Y / z + cy).astype(np.int32)

    out = np.zeros((h, w, 3), np.uint8)
    zbuf = np.full((h, w), np.inf, np.float32)
    cols = rgb[ys, xs]
    order = np.argsort(-z)  # far first, so near overwrite
    u, v, z, cols = u[order], v[order], z[order], cols[order]

    for dy in range(-radius, radius + 1):
        for dx in range(-radius, radius + 1):
            uu, vv = u + dx, v + dy
            ok = (uu >= 0) & (uu < w) & (vv >= 0) & (vv < h)
            uu, vv, zz, cc = uu[ok], vv[ok], z[ok], cols[ok]
            closer = zz < zbuf[vv, uu]
            zbuf[vv[closer], uu[closer]] = zz[closer]
            out[vv[closer], uu[closer]] = cc[closer]
    return out[:, :, ::-1].copy(), np.isfinite(zbuf)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--clip-dir", type=Path, required=True,
                    help="outputs/moge/sample_bucket/<uuid> (needs per_frame/*.npz + stats.json)")
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--radius", type=int, default=1, help="splat half-size in px")
    ap.add_argument("--fps", type=float, default=15.0)
    args = ap.parse_args()

    stats = json.loads((args.clip_dir / "stats.json").read_text())
    video = REPO_ROOT / stats["video"]
    files = sorted((args.clip_dir / "per_frame").glob("frame_*.npz"))
    if not files:
        raise SystemExit(f"no npz under {args.clip_dir/'per_frame'}")

    cap = cv2.VideoCapture(str(video))
    rgbs = []
    while True:
        ok, f = cap.read()
        if not ok:
            break
        rgbs.append(cv2.cvtColor(f, cv2.COLOR_BGR2RGB))
    cap.release()
    n = min(len(rgbs), len(files))
    print(f"{args.clip_dir.name}: {n} frames, {stats['width']}x{stats['height']}")

    args.out.mkdir(parents=True, exist_ok=True)
    fill_fracs = []

    def frames():
        for i in range(n):
            d = np.load(files[i])
            depth = d["depth"].astype(np.float32)
            depth = np.where(d["mask"].astype(bool), depth, np.nan)
            K = d["intrinsics"].astype(np.float64).copy()
            # MoGe returns NORMALISED intrinsics (fx, fy, cx, cy in [0,1] units of
            # image size). Denormalise before any pixel arithmetic -- using them raw
            # would put every point within one pixel of the origin.
            h, w = depth.shape
            K[0, :] *= w
            K[1, :] *= h
            img, filled = splat(depth, rgbs[i], K, args.radius)
            fill_fracs.append(float(filled.mean()))
            yield img

    out_video = args.out / "control_moge_pointcloud.mp4"
    write_h264(frames(), out_video, args.fps)
    ref = args.out / "ref_frame0.png"
    cv2.imwrite(str(ref), rgbs[0][:, :, ::-1])

    meta = {"uuid": args.clip_dir.name, "source_video": stats["video"],
            "n_frames": n, "splat_radius": args.radius,
            "fill_fraction_mean": round(float(np.mean(fill_fracs)), 4),
            "fill_fraction_min": round(float(np.min(fill_fracs)), 4),
            "control": str(out_video), "ref_image": str(ref),
            "instruction": stats.get("instruction")}
    (args.out / "control_meta.json").write_text(json.dumps(meta, indent=2))
    print(f"fill fraction mean {meta['fill_fraction_mean']:.4f} "
          f"min {meta['fill_fraction_min']:.4f}")
    print(f"wrote {out_video}\nwrote {ref}")


if __name__ == "__main__":
    main()