scripts/compare_zeroshot.py

#!/usr/bin/env python
"""Contact sheet + side-by-side video for a zero-shot VACE run against its own window.

Simpler than ``compare_wan_variants.py`` on purpose. That script exists because the
A/B/C/D ablation compared generated clips against the *raw* 127-frame DROID capture,
so it had to reproduce ``VaceVideoProcessor``'s own 81-of-127 frame selection and its
resize+crop to line the two up. Here the ground truth is the window's own
``target.mp4``, which S8's exporter already wrote at exactly (81, 480, 832, 3) --
the same contract the generated clip comes out at -- so frame i corresponds to
frame i and no matching step is needed or wanted.

Metrics are deliberately secondary here. This project's own S8 note recorded that a
photoreal *hallucinated human hand* outscores a correctly-rendered robot arm on
PSNR/SSIM/LPIPS by matching the real clip's skin-tone texture and motion statistics
-- so for any "is the agent right?" question the contact sheet is the evidence and
the table is not. The numbers are printed anyway because they are cheap and do track
the separate question of overall photorealism.

Usage:
    PYTHONPATH=src python scripts/compare_zeroshot.py \\
        --generated outputs/zeroshot/zs_ours_robotprompt.mp4 \\
        --window outputs/datagen/<uuid>/<cam>/vace/window_00000_00081 \\
        --out-dir outputs/zeroshot
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import imageio.v3 as iio
import numpy as np

#: Frames sampled for the contact sheet. Same four the S8 note used, so its
#: rows and this script's rows are visually comparable at a glance.
SHEET_FRAMES = (5, 27, 50, 75)


def read_video(path: Path) -> np.ndarray:
    v = np.asarray(iio.imread(path, plugin="pyav"))
    if v.ndim != 4 or v.shape[-1] != 3:
        raise ValueError(f"{path}: expected (T,H,W,3), got {v.shape}")
    return v


def psnr(a: np.ndarray, b: np.ndarray) -> float:
    mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
    return float("inf") if mse == 0 else 20.0 * np.log10(255.0 / np.sqrt(mse))


def temporal_activity(v: np.ndarray) -> float:
    """Mean abs frame-to-frame grayscale diff on a [0,1] scale.

    The S8/ablation notes' own "is the generation actually moving, or is it a
    near-static clip that flatters every similarity metric?" check.
    """
    g = v.astype(np.float32).mean(axis=-1) / 255.0
    return float(np.abs(np.diff(g, axis=0)).mean())


def label(img: np.ndarray, text: str) -> np.ndarray:
    """Burn a text label into the top-left of a frame copy (cv2, no font files)."""
    import cv2

    out = img.copy()
    cv2.rectangle(out, (0, 0), (len(text) * 11 + 10, 26), (0, 0, 0), -1)
    cv2.putText(out, text, (5, 19), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA)
    return out


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--generated", required=True, type=Path)
    p.add_argument("--window", required=True, type=Path)
    p.add_argument("--out-dir", required=True, type=Path)
    p.add_argument("--tag", default="zeroshot")
    args = p.parse_args()

    gen = read_video(args.generated)
    real = read_video(args.window / "target.mp4")
    ctrl = read_video(args.window / "control_normal.mp4")
    caption = (args.window / "caption.txt").read_text().strip()

    n = min(len(gen), len(real), len(ctrl))
    if not (len(gen) == len(real) == len(ctrl)):
        print(f"WARNING: frame counts differ (gen={len(gen)} real={len(real)} "
              f"ctrl={len(ctrl)}); comparing the first {n}")
    gen, real, ctrl = gen[:n], real[:n], ctrl[:n]

    args.out_dir.mkdir(parents=True, exist_ok=True)

    # Contact sheet: rows = real / control / generated, cols = SHEET_FRAMES.
    frames = [f for f in SHEET_FRAMES if f < n]
    rows = []
    for name, vid in (("real", real), ("control_normal", ctrl), ("generated", gen)):
        rows.append(np.concatenate([label(vid[f], f"{name} @{f}") for f in frames], axis=1))
    sheet = np.concatenate(rows, axis=0)
    sheet_path = args.out_dir / f"{args.tag}_contact_sheet.png"
    iio.imwrite(sheet_path, sheet)

    # Side-by-side video, real on the left.
    sbs_path = args.out_dir / f"{args.tag}_side_by_side.mp4"
    iio.imwrite(
        sbs_path,
        np.concatenate([real, gen], axis=2),
        fps=16,
        codec="libx264",
        plugin="pyav",
        in_pixel_format="rgb24",
    )

    metrics = {
        "caption": caption,
        "n_frames": int(n),
        "psnr_generated_vs_real": round(psnr(gen, real), 3),
        "temporal_activity": {
            "real": round(temporal_activity(real), 5),
            "generated": round(temporal_activity(gen), 5),
            "control": round(temporal_activity(ctrl), 5),
        },
        "note": (
            "PSNR here tracks overall photorealism only. Per this project's own S8 "
            "finding, a hallucinated human hand outscores a correct robot arm on "
            "PSNR/SSIM/LPIPS -- read the contact sheet for agent identity."
        ),
    }
    metrics["temporal_activity"]["generated_vs_real_ratio"] = round(
        metrics["temporal_activity"]["generated"] / metrics["temporal_activity"]["real"], 3
    )
    (args.out_dir / f"{args.tag}_metrics.json").write_text(json.dumps(metrics, indent=2))

    print(json.dumps(metrics, indent=2))
    print(f"\ncontact sheet: {sheet_path}")
    print(f"side by side:  {sbs_path}")


if __name__ == "__main__":
    main()