scripts/render_background_probe.py

#!/usr/bin/env python
"""Does the appearance LoRA generate a *scene*, or only repaint a robot onto one?

Every appearance-pair control ever built (``scripts/appearance_control.py``, 800
clips, the whole ``outputs/appearance_pairs`` training set) is the real camera
frame with only the robot region overwritten by the URDF render
(``appearance_control.render_and_composite``, ``out = frame_bgr.copy(); out[win]
= rr.color[win]``). The model is handed the table, the drawer and the brick as
real pixels on every single training example -- it has never once had to
generate them. Every PSNR number this project has reported for the fine-tune
(e.g. the real-control run: 19.05 dB full-frame, 11.46 dB robot-region for
``step-500``, recomputed below rather than trusted) is consistent with two very
different stories: "the fine-tune learned to render this scene" or "the
fine-tune learned to paint a robot and the background score is free, because it
was never asked to produce the background." Nothing measured so far
distinguishes them.

This script does: build three controls that are IDENTICAL to the real control
in every respect except what fills the non-robot region --

  * ``plate`` -- a temporal-median background plate of this episode's own real
    video (``fpgm.datagen.plate.compute_background_plate``, reused verbatim, not
    reimplemented -- it already does exactly this and warns if the camera turns
    out not to be static, which would make the median a meaningless blur rather
    than a photograph; that warning is surfaced here, not swallowed).
  * ``grey`` -- flat mid-grey (128,128,128).
  * ``black`` -- flat black (0,0,0).

-- then generates from each with the same LoRA, same settings as every prior
run on this clip, and looks at what shows up in the non-robot region of the
output. If the model reconstructs a plausible table/drawer/brick, the
fine-tune's appearance knowledge genuinely transfers past the robot. If the
background comes out empty, smeared, or a static copy of the reference frame,
the model only ever learned "paint a robot here" and the background quality
this project has been reporting was supplied by the control, not produced by
the model.

**The reference image stays the real ``target_frame0``** (the training/
deployment contract ``sample_appearance_lora.py --reference target_frame0``
already defaults to) -- only the *control*'s background changes. This isolates
one variable: what the model does with the per-frame conditioning video, not
what it copies from the one-shot reference channel (a separate, already-measured
effect: swapping the reference alone moves robot-region PSNR by 8.8 dB while
swapping the control alone moves it 0.3 dB, per ``sample_appearance_lora.py``'s
own docstring).

--- What changed in appearance_control.py, and why so little ------------------
``render_and_composite`` gained one keyword-only parameter, ``background_bgr``.
Default ``None`` reproduces ``out = frame_bgr.copy()`` byte for byte -- the 800
already-built pairs and every existing caller (``ClipBuilder.render_loop``,
``render_counterfactual_grip.py``) pass nothing for it and are unaffected. The
alternative -- a second composite function, or a background-swap wrapped around
the existing one -- was rejected: the robot render, the depth-alignment fit, the
occlusion test and the ``win`` mask are all identical regardless of background
(none of them read ``frame_bgr`` except as the paint base), so forking the
function would duplicate exactly the logic that must NOT drift between the real
pipeline and this probe.

--- Verification, not assumption -----------------------------------------------
The claim "robot render/pose/occlusion/crop/resolution are identical across all
three variants and the real control" is checked, not asserted: every frame is
rendered independently for the ``real``/``plate``/``grey``/``black`` background
(four separate calls to ``render_and_composite``, each carrying its own
depth-scale state), and the resulting ``drawn & ~occluded`` ("win") masks are
compared pixel-for-pixel across all four. Since none of the masking math reads
the background, they are expected to match exactly (mismatch fraction should
read 0); the number is reported rather than presumed. The ``real`` pass's own
composited pixels are additionally diffed (mean abs error, robot region only)
against the existing ``outputs/appearance_holdout/.../control.mp4`` this
episode's real pair was already built with -- a nonzero number there is expected
and is H.264 compression noise (the comparison is an uncompressed render array
against a decoded lossy video), not evidence of a behaviour change; only a mask
mismatch would indicate one.

--- PSNR methodology ------------------------------------------------------------
Aggregate-MSE PSNR (10*log10(255^2 / mean squared error over every included
pixel across all frames), not a per-frame PSNR averaged afterward -- the same
convention ``scripts/compare_wan_variants.py`` uses, for the same reason: a
handful of frames with near-zero error would otherwise dominate a naive mean of
dB values. "Robot region" is the ``win`` mask actually painted into the control
(``drawn & ~occluded``); "background region" is its complement. Both real- and
synthetic-background variants share the SAME per-frame masks (verified above),
so the split is apples-to-apples across all four generations.

--- The floor ---------------------------------------------------------------
DiffSynth generation is measured NOT bit-reproducible across process launches,
even at a fixed seed on an identical control (MAE ~2.0 between two launches of
the same checkpoint) -- any PSNR difference below what that floor would produce
is not attributable to the background swap.

--- Pipeline reuse ------------------------------------------------------------
Depth is stage 1 of the existing pipeline (``appearance_depth.py``, env
``moge``) and must already exist in ``--work``; this script does not regenerate
it (the exact artifacts for this episode/camera are already on disk and are
reused as-is, per the task's instruction not to re-run stage 1). Rendering
reuses ``appearance_control.ClipBuilder.load_clip`` and
``appearance_control.render_and_composite`` unmodified except for the one new
keyword argument. Generation shells out to the unmodified
``sample_appearance_lora.py`` (env ``wan-train``) once per background variant,
exactly the two-conda-env split the rest of this pipeline already uses.

Usage (three stages: fpgm render+probe, then wan-train generate x3, orchestrated
here; only stage 1's depth is assumed already built)::

    PYOPENGL_PLATFORM=egl PYTHONPATH=src /home/quang/miniconda3/envs/fpgm/bin/python \\
        scripts/render_background_probe.py --gpu <gpu>
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import warnings
from contextlib import nullcontext
from pathlib import Path
from typing import TYPE_CHECKING

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))
sys.path.insert(0, str(REPO_ROOT / "scripts"))

if TYPE_CHECKING:
    import numpy as np

    from fpgm.utils.timing import StepTimer

_DEFAULT_UUID = "AUTOLab+0d4edc83+2023-10-21-19h-37m-47s"
_DEFAULT_SERIAL = "22008760"
_WAN_PY = "/home/quang/miniconda3/envs/wan-train/bin/python"
_GREY_VALUE = 128  # flat mid-grey; 0-255, all channels
_VARIANTS = ("plate", "grey", "black")  # the three generated controls
_ALL_PASSES = ("real", *_VARIANTS)  # + the real-frame background, for verification only


def _step(timer: StepTimer | None, label: str, **kwargs):
    """``timer.step(label, **kwargs)``, or a no-op context when ``timer is None``."""
    return timer.step(label, **kwargs) if timer is not None else nullcontext()


def build_backgrounds(
    mp4: Path, out_w: int, out_h: int, timer: StepTimer | None = None,
) -> tuple[dict, dict, list[str]]:
    """The three non-robot backgrounds, cover-cropped to the render resolution.

    Returns:
        ``(backgrounds, plate_stats, warnings_fired)`` -- ``backgrounds`` maps
        ``{"plate", "grey", "black"} -> (out_h, out_w, 3)`` uint8 BGR arrays,
        ``plate_stats`` is ``fpgm.datagen.plate.PlateStats.__dict__`` for the
        median build, and ``warnings_fired`` lists any
        ``NonStaticCameraWarning`` messages -- empty when the camera checked out
        static, non-empty (and MUST be read, not ignored) otherwise, since a
        non-static camera makes the plate a meaningless blur rather than a
        background photograph (see ``fpgm.datagen.plate`` module docstring).
    """
    from dataclasses import asdict

    import appearance_control as ac
    import cv2
    import numpy as np

    from fpgm.datagen.plate import NonStaticCameraWarning, compute_background_plate

    with _step(timer, "build_plate"):
        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always", NonStaticCameraWarning)
            plate_rgb, stats = compute_background_plate(mp4)
        fired = [str(w.message) for w in caught if issubclass(w.category, NonStaticCameraWarning)]
        plate_bgr = ac.cover_crop(cv2.cvtColor(plate_rgb, cv2.COLOR_RGB2BGR), out_w, out_h)

    backgrounds = {
        "plate": plate_bgr,
        "grey": np.full((out_h, out_w, 3), _GREY_VALUE, np.uint8),
        "black": np.zeros((out_h, out_w, 3), np.uint8),
    }
    return backgrounds, asdict(stats), fired


def render_variants(
    builder, clip: dict, depth_all: np.ndarray, backgrounds: dict,
    real_control_mp4: Path | None, out_root: Path, fps: float,
    timer: StepTimer | None = None,
) -> tuple[dict, np.ndarray]:
    """One synchronized pass over the episode: 4 independent composites per frame.

    Renders ``real`` (background_bgr=None, i.e. the unmodified default path) and
    the three synthetic-background variants for every frame, from the SAME
    joint pose / gripper / camera / MoGe depth -- only ``background_bgr``
    differs between the four ``render_and_composite`` calls. Writes
    ``control.mp4`` for ``plate``/``grey``/``black`` into
    ``out_root/pairs/<variant>/``; does not write a ``real`` video (that
    already exists at ``outputs/appearance_holdout/...``) but still runs the
    ``real`` pass, purely to verify against it -- see module docstring.

    Returns:
        ``(stats, win_masks)`` -- ``stats`` has the mask-agreement and (if
        ``real_control_mp4`` given) pixel-diff numbers; ``win_masks`` is
        ``(n, out_h, out_w)`` uint8, the real pass's own ``drawn & ~occluded``
        mask per frame -- the robot-region definition every downstream PSNR
        split uses, saved once here because the three synthetic variants were
        verified (not assumed) to share it.
    """
    import appearance_control as ac
    import cv2
    import numpy as np

    joint_positions, gripper = clip["joint_positions"], clip["gripper"]
    mp4, camera = clip["mp4"], clip["camera"]
    frames_per_step, n = clip["frames_per_step"], clip["n"]
    out_w, out_h = clip["out_w"], clip["out_h"]
    n_rows = joint_positions.shape[0]
    rend = builder.renderer(out_w, out_h)

    writers = {}
    for v in _VARIANTS:
        d = out_root / "pairs" / v
        d.mkdir(parents=True, exist_ok=True)
        writers[v] = ac.H264Writer(d / "control.mp4", out_w, out_h, fps, builder.crf)

    last_a = dict.fromkeys(_ALL_PASSES, 1.0)
    win_masks = np.zeros((n, out_h, out_w), np.uint8)
    total_drawn = 0
    mask_mismatch = 0
    mae_sum, mae_n = 0.0, 0

    cap = cv2.VideoCapture(str(mp4))
    real_cap = cv2.VideoCapture(str(real_control_mp4)) if real_control_mp4 else None
    try:
        with _step(timer, "render_probe", n=n):
            for t in range(n):
                ok, frame_bgr = cap.read()
                if not ok:
                    break
                frame_bgr = ac.cover_crop(frame_bgr, out_w, out_h)
                row = min(int(round(t / frames_per_step)), n_rows - 1)
                moge = cv2.resize(depth_all[t].astype(np.float32), (out_w, out_h),
                                  interpolation=cv2.INTER_NEAREST)

                passes = {}
                for v in _ALL_PASSES:
                    bg = None if v == "real" else backgrounds[v]
                    out, drawn, occluded, a, _resid = ac.render_and_composite(
                        rend, builder.robot, joint_positions[row], gripper[row],
                        frame_bgr, moge, camera, builder.depth_tol_m, last_a[v],
                        background_bgr=bg)
                    last_a[v] = a
                    passes[v] = (out, drawn & ~occluded)

                win_real = passes["real"][1]
                win_masks[t] = win_real.astype(np.uint8)
                total_drawn += int(win_real.sum())
                for v in _VARIANTS:
                    win_v = passes[v][1]
                    mask_mismatch += int(np.count_nonzero(win_v != win_real))
                    writers[v].write(passes[v][0])

                if real_cap is not None:
                    ok2, real_ctrl_frame = real_cap.read()
                    if ok2 and win_real.any():
                        diff = np.abs(passes["real"][0][win_real].astype(np.float64)
                                      - real_ctrl_frame[win_real].astype(np.float64))
                        mae_sum += float(diff.sum())
                        mae_n += diff.size
    finally:
        cap.release()
        if real_cap is not None:
            real_cap.release()
        for w in writers.values():
            w.close()

    stats = {
        "n_frames": n,
        "total_drawn_px": total_drawn,
        "mask_mismatch_px": mask_mismatch,
        "mask_mismatch_fraction": mask_mismatch / max(total_drawn, 1),
        "real_control_pixel_mae": (mae_sum / mae_n) if mae_n else None,
        "real_control_pixel_mae_note": (
            "mean abs diff, robot region only, uncompressed render vs H.264-decoded "
            "outputs/appearance_holdout .../control.mp4 -- nonzero is expected "
            "compression noise, not a behaviour change; only mask_mismatch_fraction "
            "checks the algorithm itself"
        ),
    }
    return stats, win_masks


def generate_variant(
    pair_dir: Path, samples_dir: Path, lora: Path, *, gpu: int, seed: int, steps: int,
    cfg_scale: float, num_frames: int, chunk_frames: int, baseline: bool,
    offload_text_encoder: bool, timer: StepTimer | None = None, label: str = "",
) -> None:
    """Shell out to the unmodified ``sample_appearance_lora.py`` (env ``wan-train``)."""
    cmd = [
        _WAN_PY, str(REPO_ROOT / "scripts/sample_appearance_lora.py"),
        "--pair", str(pair_dir), "--out", str(samples_dir), "--lora", str(lora),
        "--start", "0", "--num-frames", str(num_frames),
        "--chunk-frames", str(chunk_frames), "--chain",
        "--seed", str(seed), "--steps", str(steps), "--cfg-scale", str(cfg_scale),
    ]
    if baseline:
        cmd.append("--baseline")
    if offload_text_encoder:
        cmd.append("--offload-text-encoder")
    env = {**os.environ, "CUDA_VISIBLE_DEVICES": str(gpu)}
    with _step(timer, f"generate_{label}", n=num_frames * (2 if baseline else 1)):
        print(f"running ({label}):", " ".join(cmd), flush=True)
        result = subprocess.run(cmd, env=env, cwd=REPO_ROOT)
    if result.returncode != 0:
        raise SystemExit(f"sample_appearance_lora.py failed for {label} (exit {result.returncode})")


def masked_psnr(gen: np.ndarray, real: np.ndarray, mask: np.ndarray) -> float:
    """Aggregate-MSE PSNR over ``mask``-selected pixels. See module docstring."""
    import numpy as np

    if not mask.any():
        return float("nan")
    se = (gen[mask].astype(np.float64) - real[mask].astype(np.float64)) ** 2
    mse = float(se.mean())
    return float("inf") if mse == 0.0 else float(10.0 * np.log10(255.0 ** 2 / mse))


def read_frames_bgr(path: Path, n: int) -> np.ndarray:
    import cv2
    import numpy as np

    cap = cv2.VideoCapture(str(path))
    frames = []
    try:
        for _ in range(n):
            ok, bgr = cap.read()
            if not ok:
                break
            frames.append(bgr)
    finally:
        cap.release()
    return np.stack(frames)


def compute_metrics(
    win_masks: np.ndarray, target_path: Path, generated: dict[str, Path],
    n_frames: int, timer: StepTimer | None = None,
) -> dict:
    """Full-frame and background-region (non-robot) PSNR for every generated clip.

    ``win_masks`` (from :func:`render_variants`) is truncated to ``n_frames``
    (the generation window is shorter than the full episode -- 249 of 251, so
    the tail two frames' masks are simply unused). ``target_path`` is one
    ``target.mp4`` shared by every generation (the real frame window is
    identical regardless of which control produced the output).
    """
    import numpy as np

    with _step(timer, "compute_metrics", n=len(generated)):
        target = read_frames_bgr(target_path, n_frames)
        masks = win_masks[:n_frames].astype(bool)
        full_mask = np.ones_like(masks)
        results = {}
        for label, path in generated.items():
            if not path.exists():
                results[label] = {"error": f"missing: {path}"}
                continue
            gen = read_frames_bgr(path, n_frames)
            k = min(gen.shape[0], target.shape[0], masks.shape[0])
            results[label] = {
                "n_frames": k,
                "psnr_full_frame_db": masked_psnr(gen[:k], target[:k], full_mask[:k]),
                "psnr_background_db": masked_psnr(gen[:k], target[:k], ~masks[:k]),
                "psnr_robot_region_db": masked_psnr(gen[:k], target[:k], masks[:k]),
            }
    return results


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--episode", default=_DEFAULT_UUID)
    ap.add_argument("--camera", default=_DEFAULT_SERIAL)
    ap.add_argument("--work", type=Path, default=REPO_ROOT / "outputs/counterfactual_grip/work",
                    help="MoGe depth artifacts (stage 1); must already exist, not re-run here")
    ap.add_argument("--out", type=Path, default=REPO_ROOT / "outputs/background_probe")
    ap.add_argument("--real-pair", type=Path,
                    default=REPO_ROOT / "outputs/appearance_holdout"
                                        f"/{_DEFAULT_UUID}__{_DEFAULT_SERIAL}",
                    help="existing real-frame-background pair, for caption/meta reuse "
                         "and for the render verification's pixel-diff check")
    ap.add_argument("--real-samples-dir", type=Path,
                    default=REPO_ROOT / "outputs/lora_samples_rigid",
                    help="existing real-control step-500 generation (same seed/steps/"
                         "cfg/chunking) to recompute the real-control reference PSNR "
                         "numbers from, instead of re-generating them")
    ap.add_argument("--episodes-root", type=Path, default=REPO_ROOT / "data/droid_raw")
    ap.add_argument("--cameras-dir", type=Path,
                    default=REPO_ROOT / "data/pointworld/droid/cameras")
    ap.add_argument("--intrinsics", type=Path,
                    default=REPO_ROOT / "configs/droid_camera_intrinsics.json")
    ap.add_argument("--urdf", type=Path, default=None)
    ap.add_argument("--depth-tol-m", type=float, default=0.15)
    ap.add_argument("--crf", type=int, default=18)
    ap.add_argument("--fps", type=float, default=15.0)
    ap.add_argument("--gpu", type=int, default=0,
                    help="sets EGL_DEVICE_ID for rendering and CUDA_VISIBLE_DEVICES "
                         "for each wan-train generation subprocess")
    ap.add_argument("--lora", type=Path,
                    default=REPO_ROOT / "outputs/lora_appearance_2k/step-500.safetensors")
    ap.add_argument("--num-frames", type=int, default=249)
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--cfg-scale", type=float, default=5.0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--chunk-frames", type=int, default=81)
    ap.add_argument("--offload-text-encoder", action="store_true")
    ap.add_argument("--skip-generate", action="store_true",
                    help="render the three controls and verify, then stop -- for "
                         "iterating on the render/verification without paying for "
                         "a wan-train generation each time")
    ap.add_argument("--skip-metrics", action="store_true",
                    help="skip the PSNR pass (e.g. generation already ran and this "
                         "is a render-only re-check)")
    return ap.parse_args()


def main() -> int:
    args = parse_args()
    os.environ["EGL_DEVICE_ID"] = str(args.gpu)
    os.environ.setdefault("PYOPENGL_PLATFORM", "egl")

    import appearance_control as ac
    import numpy as np

    from fpgm.utils.timing import StepTimer

    timer = StepTimer("background_probe")
    args.out.mkdir(parents=True, exist_ok=True)

    depth_json = args.work / f"{args.episode}__{args.camera}.depth.json"
    depth_npy = args.work / f"{args.episode}__{args.camera}.depth.npy"
    if not (depth_json.exists() and depth_npy.exists()):
        raise SystemExit(
            f"missing MoGe depth for {args.episode}/{args.camera} in {args.work} -- "
            f"regenerate with scripts/appearance_depth.py (env moge) first"
        )
    depth_meta = json.loads(depth_json.read_text())

    urdf = args.urdf
    if urdf is None:
        from fpgm.config_datagen import DatagenProfile
        urdf = Path(DatagenProfile.from_yaml(
            str(REPO_ROOT / "configs/datagen_droid.yaml")).paths.urdf)

    real_pair_meta = json.loads((args.real_pair / "meta.json").read_text())

    builder = ac.ClipBuilder(
        urdf, ac.load_intrinsic_table(args.intrinsics), args.depth_tol_m, args.crf)
    try:
        with timer.step("load_clip"):
            clip = builder.load_clip(depth_meta, args.episodes_root, args.cameras_dir)

        backgrounds, plate_stats, plate_warnings = build_backgrounds(
            clip["mp4"], clip["out_w"], clip["out_h"], timer=timer)
        if plate_warnings:
            print(f"WARNING: plate build fired NonStaticCameraWarning: {plate_warnings}",
                  flush=True)
        else:
            print(f"plate: static-camera check OK "
                  f"(corner_median_abs_diff={plate_stats['corner_median_abs_diff']:.3f})",
                  flush=True)

        depth_all = np.load(depth_npy, mmap_mode="r")
        render_stats, win_masks = render_variants(
            builder, clip, depth_all, backgrounds,
            args.real_pair / "control.mp4", args.out, args.fps, timer=timer)
        print(
            f"mask agreement: {render_stats['mask_mismatch_px']}/"
            f"{render_stats['total_drawn_px']} drawn px differ across "
            f"real/plate/grey/black ({render_stats['mask_mismatch_fraction']*100:.4f}%); "
            f"real-vs-existing-holdout robot-region pixel MAE="
            f"{render_stats['real_control_pixel_mae']}", flush=True,
        )

        with timer.step("write_pairs"):
            np.save(args.out / "robot_mask.npy", win_masks)
            for v in _VARIANTS:
                pair_dir = args.out / "pairs" / v
                shutil.copyfile(args.real_pair / "target.mp4", pair_dir / "target.mp4")
                (pair_dir / "meta.json").write_text(json.dumps({
                    "uuid": real_pair_meta["uuid"],
                    "camera_serial": real_pair_meta["camera_serial"],
                    "n_frames": render_stats["n_frames"],
                    "fps": args.fps,
                    "video_wh": [clip["out_w"], clip["out_h"]],
                    "caption": real_pair_meta["caption"],
                    "lab": real_pair_meta.get("lab"),
                    "control": "control.mp4",
                    "target": "target.mp4",
                    "background": v,
                    "background_probe_note": (
                        "non-robot region is a synthetic background, not the real "
                        "camera frame -- see scripts/render_background_probe.py"
                    ),
                }, indent=2))
            verification = dict(render_stats)
            verification["plate_stats"] = plate_stats
            verification["plate_non_static_camera_warnings"] = plate_warnings
            (args.out / "verification.json").write_text(json.dumps(verification, indent=2))
    finally:
        builder.close()

    print(f"wrote {args.out / 'pairs'} (plate/grey/black), robot_mask.npy, "
          f"verification.json", flush=True)

    if args.skip_generate:
        print(timer.report(), flush=True)
        return 0

    samples_root = args.out / "samples"
    for v in _VARIANTS:
        generate_variant(
            args.out / "pairs" / v, samples_root / v, args.lora, gpu=args.gpu,
            seed=args.seed, steps=args.steps, cfg_scale=args.cfg_scale,
            num_frames=args.num_frames, chunk_frames=args.chunk_frames,
            baseline=(v == "plate"), offload_text_encoder=args.offload_text_encoder,
            timer=timer, label=v,
        )

    if args.skip_metrics:
        print(timer.report(), flush=True)
        return 0

    tag = f"{args.episode}__{args.camera}_f00000_chain{args.chunk_frames}"
    generated = {
        "plate_baseline": samples_root / "plate" / f"{tag}__baseline.mp4",
        "plate_lora": samples_root / "plate" / f"{tag}__lora_{args.lora.stem}.mp4",
        "grey_lora": samples_root / "grey" / f"{tag}__lora_{args.lora.stem}.mp4",
        "black_lora": samples_root / "black" / f"{tag}__lora_{args.lora.stem}.mp4",
        "real_lora": args.real_samples_dir / f"{tag}__lora_{args.lora.stem}.mp4",
    }
    real_baseline = args.real_samples_dir / f"{tag}__baseline.mp4"
    if real_baseline.exists():
        generated["real_baseline"] = real_baseline

    metrics = compute_metrics(
        win_masks, samples_root / "plate" / f"{tag}__target.mp4",
        generated, args.num_frames, timer=timer)
    (args.out / "metrics.json").write_text(json.dumps(metrics, indent=2))
    print("\nPSNR (dB), aggregate-MSE, robot-region mask shared across all rows:", flush=True)
    for label, m in metrics.items():
        if "error" in m:
            print(f"  {label}: {m['error']}", flush=True)
        else:
            print(f"  {label:16s} full={m['psnr_full_frame_db']:6.2f}  "
                  f"background={m['psnr_background_db']:6.2f}  "
                  f"robot_region={m['psnr_robot_region_db']:6.2f}", flush=True)
    print(f"wrote {args.out / 'metrics.json'}", flush=True)

    print(timer.report(), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())