scripts/export_zeroshot_control.py

#!/usr/bin/env python
"""Write the trainer's own composited RGB control video for one window, standalone.

Exists because a zero-shot probe of the pretrained VACE checkpoint must be fed
*the same tensor the trainer will feed it*, not a proxy. The first probe run in
this project passed ``control_depth.mkv`` alone and produced a false-colour
depth passthrough -- but the trainer never sees that file on its own:
``fpgm.training.types.BundleAssemblyConfig`` records that VACE's branch has one
``vace_patch_embedding`` taking one ``vace_video``, so the three control
channels are composited into a single RGB video (R=depth, G=seg id normalised,
B=normal-Z) before they reach it. That probe therefore tested a strictly poorer
input than training uses, and its result cannot be carried over.

This reuses :func:`fpgm.training.bundle_io.load_window_arrays` -- the trainer's
actual decode+composite path -- rather than reimplementing the channel order or
the seg normalisation, so the probe cannot silently drift from the trainer.

Written as FFV1-in-Matroska (lossless). H.264 would be the obvious choice for
something a video model reads, but the G channel is *segmentation ids*, and
:class:`fpgm.viz.video.LosslessWriter`'s own docstring records why that is
unsafe: lossy ringing shifting an id by +-1 relabels robot pixels as object
pixels. decord decodes FFV1/mkv fine here (verified against this window's own
``control_depth.mkv``).

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

from __future__ import annotations

import argparse
import subprocess
from pathlib import Path

import numpy as np

from fpgm.training.bundle_io import load_window_arrays
from fpgm.training.manifest import _parse_window_dir
from fpgm.training.types import BundleAssemblyConfig
from fpgm.viz.video import _resolve_ffmpeg


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.

    ``fpgm.viz.video.LosslessWriter`` is grayscale-only by construction, so this
    uses the same codec/container through its own pipe rather than widening a
    class whose docstring is explicitly about single-channel integer data.
    """
    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)
    args = ap.parse_args()

    sample = _parse_window_dir(args.window)
    cfg = BundleAssemblyConfig()
    arrays = load_window_arrays(sample, cfg)
    control = arrays["control"]

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

    per_channel = control.reshape(-1, 3).mean(axis=0).round(2)
    print(f"window:      {args.window}")
    print(f"control:     {control.shape} {control.dtype}  (R=depth, G=seg/{cfg.max_seg_id}, B=normal-Z)")
    print(f"channel mean R/G/B: {per_channel}")
    print(f"nonzero seg pixels: {(control[..., 1] > 0).mean():.4%}")
    print(f"caption:     {arrays['caption']}")
    print(f"wrote:       {args.out}")


if __name__ == "__main__":
    main()