scripts/export_vis_control.py

#!/usr/bin/env python
"""Build a Cosmos-``vis``-style control video: shaded CG foreground over the real plate.

**Why this exists.** Two zero-shot probes of the pretrained VACE-1.3B checkpoint
on this project's own control buffers -- ``control_depth.mkv`` alone, then the
trainer's full R=depth/G=seg/B=normal-Z composite -- both produced a false-colour
passthrough of the control rather than a photograph (PSNR 9.93 / 9.17, generated
temporal activity 0.00586 / 0.00564 against a control's 0.00531 and a real
0.01311). The mechanism is visible in ``third_party/wan2.1/wan/vace.py``:
``vace_encode_frames`` runs ``vae.encode(frames)`` on the control video, so the
control passes through Wan's own VAE -- trained on natural video -- before the
DiT ever sees it. A false-colour geometry buffer is out of distribution at the
*encoder*, not merely at the denoiser.

Cosmos-Transfer2.5's ``vis`` control does not have this problem because its
control already *is* a photograph: ``outputs/cosmos_input/input_video.mp4`` is
the CG foreground composited over the temporal-median plate, and those runs came
out photoreal. This script applies the same idea to VACE's single control slot.

**What the foreground shading is, stated plainly.** The datagen ``vace/`` bundles
carry no RGB render -- only depth, seg ids and normals -- so the foreground here
is a Lambertian shade computed from the window's own decoded normals, not a
materials render. It is a grey CG robot, in the spirit of Cosmos's own CG-over-
plate input, not a photoreal one. If this probe works, the honest next step is
the real renderer, not this approximation.

**Alignment.** ``ref_plate.png`` is written at native 1280x720 and deliberately
does *not* go through the exporter's ``_CoverCrop`` (see that class's docstring:
the plate is left for VACE's own letterbox-on-white reference preprocessing).
Every control buffer *does*. So the plate must be pushed through the identical
``_CoverCrop`` here before compositing, or the background lands a few percent off
the geometry it is supposed to sit behind. The exporter's own class is imported
rather than reimplemented for exactly that reason.

Usage:
    PYTHONPATH=src python scripts/export_vis_control.py \\
        --window outputs/datagen/<uuid>/<cam>/vace/window_00000_00081 \\
        --out outputs/zeroshot/control_vis.mkv
"""

from __future__ import annotations

import argparse
import subprocess
from pathlib import Path

import cv2
import numpy as np

from fpgm.datagen.export_vace import _CoverCrop
from fpgm.geometry.normals import decode_normals_rgb
from fpgm.training.bundle_io import read_id_video, read_rgb_video
from fpgm.training.manifest import _parse_window_dir
from fpgm.training.types import BundleAssemblyConfig
from fpgm.viz.video import _resolve_ffmpeg

#: Light direction in camera frame. Pointing back toward the camera and slightly
#: up-left, so the shade has some gradient rather than a flat silhouette.
#: `normals_from_depth`'s convention has a camera-facing surface at (0, 0, -1)
#: (see that function's docstring), hence the negative Z component here.
_LIGHT = np.array([-0.35, -0.45, -0.82], dtype=np.float32)

#: Lambertian ambient/diffuse split. Ambient keeps unlit faces from going pure
#: black, which would read as a hole in the plate rather than as a dark surface.
_AMBIENT = 0.35
_DIFFUSE = 0.65


def shade_from_normals(normal_rgb: np.ndarray) -> np.ndarray:
    """``(T,H,W,3)`` encoded normals -> ``(T,H,W)`` uint8 Lambertian grey.

    Decoded via :func:`fpgm.geometry.normals.decode_normals_rgb`, the exact
    inverse of the encoder ``export_vace`` used, so the shading is computed from
    the same unit vectors the exporter wrote rather than from a re-derived
    RGB<->[-1,1] mapping.
    """
    light = _LIGHT / np.linalg.norm(_LIGHT)
    out = np.empty(normal_rgb.shape[:3], dtype=np.uint8)
    for t in range(normal_rgb.shape[0]):
        n = decode_normals_rgb(normal_rgb[t])
        lam = np.clip((n * light).sum(axis=-1), 0.0, 1.0)
        out[t] = np.clip((_AMBIENT + _DIFFUSE * lam) * 255.0, 0, 255).astype(np.uint8)
    return out


def write_ffv1_rgb(frames: np.ndarray, out: Path, fps: int = 16) -> None:
    """``(T,H,W,3)`` uint8 -> lossless FFV1 in Matroska via an ffmpeg pipe."""
    t, h, w, _ = frames.shape
    out.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        _resolve_ffmpeg(), "-y", "-hide_banner", "-loglevel", "error",
        "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{w}x{h}", "-r", str(fps),
        "-i", "pipe:0", "-c:v", "ffv1", "-level", "3", "-pix_fmt", "gbrp", str(out),
    ]
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE)
    assert p.stdin is not None
    try:
        for i in range(t):
            p.stdin.write(frames[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("--window", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--fps", type=int, default=16)
    ap.add_argument("--preview", type=Path, default=None,
                    help="optional PNG of frame 27, to eyeball alignment before a GPU run")
    args = ap.parse_args()

    sample = _parse_window_dir(args.window)
    cfg = BundleAssemblyConfig()
    n = sample.frame_range[1] - sample.frame_range[0]

    seg = read_id_video(sample.control_seg, n, cfg.max_seg_id)
    normal_rgb = read_rgb_video(sample.control_normal, n)
    shade = shade_from_normals(normal_rgb)

    plate_bgr = cv2.imread(str(sample.ref_plate), cv2.IMREAD_COLOR)
    if plate_bgr is None:
        raise SystemExit(f"could not read {sample.ref_plate}")
    plate = plate_bgr[..., ::-1]

    h, w = seg.shape[1:3]
    crop = _CoverCrop.build(plate.shape[1], plate.shape[0], w, h)
    plate_aligned = crop.apply_image(plate)

    fg = seg > 0
    frames = np.repeat(plate_aligned[None], n, axis=0).copy()
    fg3 = np.repeat(fg[..., None], 3, axis=-1)
    shade3 = np.repeat(shade[..., None], 3, axis=-1)
    frames[fg3] = shade3[fg3]

    write_ffv1_rgb(frames, args.out, fps=args.fps)

    if args.preview is not None:
        args.preview.parent.mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(args.preview), frames[min(27, n - 1)][..., ::-1])

    print(f"window:        {args.window}")
    print(f"plate:         {plate.shape} -> cover-cropped to {plate_aligned.shape}")
    print(f"foreground:    {fg.mean():.4%} of pixels (seg id > 0)")
    print(f"control:       {frames.shape} {frames.dtype}  (real plate + Lambertian CG fg)")
    print(f"channel mean:  {frames.reshape(-1, 3).mean(axis=0).round(2)}")
    print(f"wrote:         {args.out}")
    if args.preview:
        print(f"preview:       {args.preview}")


if __name__ == "__main__":
    main()