scripts/compare_wan_variants.py

#!/usr/bin/env python
"""Metric + contact-sheet comparison for the Wan2.1-VACE-1.3B control-format
ablation (variants A/B/C/D, see outputs/wan_ablation/{A,B,C,D}/invocation.txt).

The real DROID clip is ground truth for all four variants:
    data/droid_raw/AUTOLab+0d4edc83+2023-10-21-19h-07m-04s/recordings/MP4/22008760.mp4

Resolution / frame-count matching (stated explicitly, see the plan's
requirement not to leave this implicit):
  - Wan2.1-VACE's own `VaceVideoProcessor` (config: vace-1.3B, size=832*480)
    deterministically selects 81 of the real video's 127 frames and a
    832x480 target crop from its own 1280x720 native size, given only
    fps/frame-count/aspect (not pixel content) -- see
    third_party/wan2.1/wan/utils/vace_processor.py
    :`_get_frameid_bbox_adjust_last` (keep_last=True at this pinned config).
    Because all four control videos share the real clip's frame count (127)
    and fps (60), this same selection applies positionally to all of them.
  - That exact selection was computed once by loading the upstream class
    directly (no reimplementation) against the real clip's actual decord
    frame timestamps -- see
    /tmp/.../scratchpad/compute_frame_match.py (run under the wan-vace conda
    env, which has decord) -- and cached into frame_match.json (copied next
    to this script's outputs at outputs/wan_ablation/frame_match.json).
    frame_ids = the 81 real-frame indices; target (oh, ow) = (480, 832).
  - The 81 matching REAL frames are extracted at those frame_ids and put
    through the IDENTICAL resize+crop transform Wan's own
    VaceVideoProcessor.resize_crop applies (scale = max(ow/iw, oh/ih),
    torchvision bicubic+antialias resize, then centre crop) -- reproduced
    here as `resize_crop_like_wan` -- to build `real_matched`: an
    (81, 480, 832, 3) uint8 sequence directly comparable, pixel-for-pixel
    position, to every generated (81, 480, 832, 3) variant.
  - The robot foreground mask (`master/robot_buffers/seg.npy`, (127,720,1280)
    uint8, 1=robot) is resampled through the SAME frame_ids + resize_crop
    (nearest-neighbor, then thresholded) to get an (81, 480, 832) bool mask
    used for the background/foreground metric split. Object masks
    (`object_*_masks.h5`) were checked and are populated at only 1/127
    frames each (see outputs/wan_ablation/C/invocation.txt) -- not usable,
    so the split is robot-vs-everything-else, not robot+objects-vs-rest.

Run with the `trellis2` conda env's python (has torch+cuda, cv2, skimage,
lpips already installed -- no new heavy deps):
    /home/quang/miniconda3/envs/trellis2/bin/python scripts/compare_wan_variants.py

CAUGHT CONFOUND (fixed before any conclusion was drawn -- recorded here so it
isn't silently lost): the first control video built for B/C was an upsample
of master/depth_preview.mp4, which was itself produced by an S3-only
dense_depth run with no robot composited in (127 near-identical frames --
see outputs/wan_ablation/B/invocation.txt for the full story and
outputs/wan_ablation/B/*_CONFOUNDED_* for the artifacts). That made the
original B/C generations a test of "a frozen control video", not "our depth
format". Fixed by compositing S2's robot_buffers/depth_mm.npy onto the
upsampled background plate (fpgm.datagen.dense_depth.compose_dense_depth)
before encoding. Because this bit us once already, `CONTROL_VIDEOS` below
computes temporal activity on every control video too, not just the
generated outputs -- if a future control is accidentally static, this table
catches it before a GPU run is spent on it.

Variant A2 (added after the fact): identical control video and seed to A,
reworded prompt (drops "LEGO" to test whether A's LEGO-styled output was a
prompt-hijack artifact rather than a control-format/model-ceiling issue).
"""
import json
import sys
from pathlib import Path

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from skimage.metrics import peak_signal_noise_ratio as sk_psnr
from skimage.metrics import structural_similarity as sk_ssim

REPO = Path("/srv/data/foundation-physics-graph-model")
OUT_DIR = REPO / "outputs/wan_ablation"
REAL_VIDEO = REPO / "data/droid_raw/AUTOLab+0d4edc83+2023-10-21-19h-07m-04s/recordings/MP4/22008760.mp4"
SEG_PATH = REPO / "outputs/datagen/AUTOLab+0d4edc83+2023-10-21-19h-07m-04s/22008760/master/robot_buffers/seg.npy"
FRAME_MATCH_JSON = OUT_DIR / "frame_match.json"

VARIANTS = ["A", "A2", "B", "C", "D"]
CONTACT_SHEET_IDX = [5, 27, 50, 75]  # of the 81 matched/generated frames

# Control video actually fed as --src_video for each variant (native
# 1280x720/127f, pre-VaceVideoProcessor). Used only to report each control's
# OWN temporal activity next to its generated output's -- catches a static
# control before/while interpreting a static output (see module docstring:
# this is exactly what the original B/C confound looked like).
CONTROL_VIDEOS = {
    "A": OUT_DIR / "A/control_depth_vace_1280x720.mp4",
    "A2": OUT_DIR / "A2/control_depth_vace_1280x720.mp4",
    "B": OUT_DIR / "B/control_depth_ours_1280x720.mp4",
    "C": OUT_DIR / "C/control_depth_ours_1280x720.mp4",
    "D": OUT_DIR / "D/control_depth_vace_1280x720.mp4",
}


def read_all_frames_rgb(path):
    cap = cv2.VideoCapture(str(path))
    frames = []
    while True:
        ret, f = cap.read()
        if not ret:
            break
        frames.append(cv2.cvtColor(f, cv2.COLOR_BGR2RGB))
    cap.release()
    return frames


def resize_crop_like_wan(img_hw_c, oh, ow, interp="bicubic"):
    """Reproduce wan/utils/vace_processor.py VaceVideoProcessor.resize_crop
    for a single HxWxC (or HxW) uint8/bool array: scale-to-cover + centre
    crop, resize via torch (antialiased bicubic for images, nearest for
    masks so the binary mask stays binary)."""
    arr = np.asarray(img_hw_c)
    single_channel = arr.ndim == 2
    if single_channel:
        arr = arr[..., None]
    ih, iw = arr.shape[:2]
    t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).float()
    if ih != oh or iw != ow:
        scale = max(ow / iw, oh / ih)
        new_h, new_w = round(scale * ih), round(scale * iw)
        if interp == "nearest":
            t = F.interpolate(t, size=(new_h, new_w), mode="nearest")
        else:
            t = F.interpolate(t, size=(new_h, new_w), mode="bicubic", antialias=True)
        x1 = (new_w - ow) // 2
        y1 = (new_h - oh) // 2
        t = t[:, :, y1:y1 + oh, x1:x1 + ow]
    out = t.squeeze(0).permute(1, 2, 0).clamp(0, 255).round().numpy().astype(np.uint8)
    if single_channel:
        out = out[..., 0]
    return out


def build_real_matched_and_mask(frame_match):
    frame_ids = frame_match["frame_ids"]
    oh, ow = frame_match["target_oh"], frame_match["target_ow"]
    real_frames = read_all_frames_rgb(REAL_VIDEO)
    seg = np.load(SEG_PATH, mmap_mode="r")
    real_matched = np.stack(
        [resize_crop_like_wan(real_frames[i], oh, ow, "bicubic") for i in frame_ids]
    )
    fg_mask = np.stack(
        [
            resize_crop_like_wan((np.asarray(seg[i]) == 1), oh, ow, "nearest")
            for i in frame_ids
        ]
    ).astype(bool)
    return real_matched, fg_mask, (oh, ow)


def temporal_activity(frames_uint8):
    """Mean absolute frame-to-frame difference, grayscale, normalized to
    [0, 1]. frames_uint8: (T, H, W, 3) uint8."""
    gray = frames_uint8.astype(np.float32).mean(axis=-1) / 255.0
    diffs = np.abs(gray[1:] - gray[:-1])
    return float(diffs.mean())


def masked_mean(per_pixel_map, mask):
    """per_pixel_map: (T,H,W). mask: (T,H,W) bool. Returns overall/bg/fg means."""
    overall = float(per_pixel_map.mean())
    bg = float(per_pixel_map[~mask].mean()) if (~mask).any() else float("nan")
    fg = float(per_pixel_map[mask].mean()) if mask.any() else float("nan")
    return overall, bg, fg


def compute_psnr_ssim(gen, real_matched, mask):
    """gen, real_matched: (T,H,W,3) uint8. mask: (T,H,W) bool (True=foreground).
    Returns dict of overall/bg/fg PSNR (dB, computed from masked MSE, not
    averaged per-frame PSNR) and SSIM (averaged per-pixel SSIM map)."""
    T = gen.shape[0]
    sq_err = (gen.astype(np.float64) - real_matched.astype(np.float64)) ** 2
    sq_err_px = sq_err.mean(axis=-1)  # (T,H,W)
    mse_overall = sq_err_px.mean()
    mse_bg = sq_err_px[~mask].mean() if (~mask).any() else float("nan")
    mse_fg = sq_err_px[mask].mean() if mask.any() else float("nan")

    def mse_to_psnr(mse):
        if mse <= 0 or np.isnan(mse):
            return float("nan") if np.isnan(mse) else float("inf")
        return 10.0 * np.log10((255.0 ** 2) / mse)

    psnr = {
        "overall": mse_to_psnr(mse_overall),
        "background": mse_to_psnr(mse_bg),
        "foreground": mse_to_psnr(mse_fg),
    }

    ssim_maps = np.empty((T,) + gen.shape[1:3], dtype=np.float64)
    for t in range(T):
        _, smap = sk_ssim(
            gen[t], real_matched[t], channel_axis=2, data_range=255, full=True
        )
        ssim_maps[t] = smap.mean(axis=-1)
    ssim_overall, ssim_bg, ssim_fg = masked_mean(ssim_maps, mask)
    ssim = {"overall": ssim_overall, "background": ssim_bg, "foreground": ssim_fg}
    return psnr, ssim


def compute_lpips(gen, real_matched, mask, lpips_model, device):
    T = gen.shape[0]
    maps = np.empty((T,) + gen.shape[1:3], dtype=np.float64)
    with torch.no_grad():
        for t in range(T):
            a = torch.from_numpy(gen[t]).permute(2, 0, 1).float().to(device) / 127.5 - 1.0
            b = torch.from_numpy(real_matched[t]).permute(2, 0, 1).float().to(device) / 127.5 - 1.0
            out = lpips_model(a.unsqueeze(0), b.unsqueeze(0))  # (1,1,H,W)
            maps[t] = out.squeeze().cpu().numpy()
    overall, bg, fg = masked_mean(maps, mask)
    return {"overall": overall, "background": bg, "foreground": fg}


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"

    if not FRAME_MATCH_JSON.exists():
        print(f"ERROR: {FRAME_MATCH_JSON} missing -- run compute_frame_match.py "
              f"(wan-vace env, needs decord) first and copy its output there.",
              file=sys.stderr)
        sys.exit(1)
    frame_match = json.loads(FRAME_MATCH_JSON.read_text())
    print(f"frame_match: {len(frame_match['frame_ids'])} frame_ids, "
          f"target=({frame_match['target_oh']},{frame_match['target_ow']})")

    real_matched, fg_mask, (oh, ow) = build_real_matched_and_mask(frame_match)
    print(f"real_matched: {real_matched.shape}, fg_mask mean fraction: {fg_mask.mean():.4f}")

    try:
        import lpips
        lpips_model = lpips.LPIPS(net="alex", spatial=True).to(device)
        have_lpips = True
    except Exception as e:
        print(f"LPIPS unavailable ({e}) -- skipping LPIPS metrics entirely, "
              f"not substituting anything else for it.", file=sys.stderr)
        lpips_model = None
        have_lpips = False

    real_temporal = temporal_activity(real_matched)
    print(f"real (matched-cadence) temporal activity: {real_temporal:.5f}")

    results = {
        "methodology": {
            "ground_truth": str(REAL_VIDEO),
            "frame_matching": "VaceVideoProcessor._get_frameid_bbox_adjust_last (upstream, unmodified), "
                               "keep_last=True, size=832x480, config vace-1.3B -- see frame_match.json",
            "n_frames_matched": len(frame_match["frame_ids"]),
            "spatial_matching": "real frames resized (scale-to-cover, bicubic+antialias) + centre-cropped "
                                 f"to {ow}x{oh}, identical to Wan's own VaceVideoProcessor.resize_crop",
            "foreground_mask": "robot only (seg.npy==1); object masks populated at 1/127 frames only, excluded",
            "lpips_available": have_lpips,
            "lpips_net": "alex (spatial=True, per-pixel map, mask-averaged)" if have_lpips else None,
        },
        "real_baseline": {
            "temporal_activity_matched_cadence": real_temporal,
        },
        "variants": {},
    }

    for v in VARIANTS:
        gen_path = OUT_DIR / v / "generated.mp4"
        if not gen_path.exists():
            print(f"variant {v}: {gen_path} missing, skipping", file=sys.stderr)
            continue
        gen_frames = read_all_frames_rgb(gen_path)
        gen = np.stack(gen_frames)
        print(f"variant {v}: generated shape {gen.shape}")
        if gen.shape[0] != real_matched.shape[0]:
            n = min(gen.shape[0], real_matched.shape[0])
            print(f"  WARNING: frame count mismatch ({gen.shape[0]} vs {real_matched.shape[0]}), "
                  f"truncating both to {n}", file=sys.stderr)
            gen_c, real_c, mask_c = gen[:n], real_matched[:n], fg_mask[:n]
        else:
            gen_c, real_c, mask_c = gen, real_matched, fg_mask
        if gen_c.shape[1:3] != real_c.shape[1:3]:
            print(f"  WARNING: spatial size mismatch {gen_c.shape[1:3]} vs {real_c.shape[1:3]}, "
                  f"resizing generated to match", file=sys.stderr)
            gen_c = np.stack([cv2.resize(f, (real_c.shape[2], real_c.shape[1])) for f in gen_c])

        psnr, ssim = compute_psnr_ssim(gen_c, real_c, mask_c)
        lp = compute_lpips(gen_c, real_c, mask_c, lpips_model, device) if have_lpips else None
        temp_act = temporal_activity(gen_c)

        control_path = CONTROL_VIDEOS.get(v)
        control_temp_act = None
        if control_path is not None and control_path.exists():
            control_frames = np.stack(read_all_frames_rgb(control_path))
            control_temp_act = temporal_activity(control_frames)

        results["variants"][v] = {
            "n_frames": int(gen_c.shape[0]),
            "resolution": f"{gen_c.shape[2]}x{gen_c.shape[1]}",
            "control_video": str(control_path) if control_path else None,
            "control_temporal_activity": control_temp_act,
            "psnr_db": psnr,
            "ssim": ssim,
            "lpips": lp,
            "temporal_activity": temp_act,
            "temporal_activity_vs_real_ratio": temp_act / real_temporal if real_temporal > 0 else float("nan"),
        }
        print(f"  PSNR overall={psnr['overall']:.2f} bg={psnr['background']:.2f} fg={psnr['foreground']:.2f}")
        print(f"  SSIM overall={ssim['overall']:.4f} bg={ssim['background']:.4f} fg={ssim['foreground']:.4f}")
        if lp:
            print(f"  LPIPS overall={lp['overall']:.4f} bg={lp['background']:.4f} fg={lp['foreground']:.4f}")
        print(f"  temporal_activity={temp_act:.5f} (real={real_temporal:.5f}, ratio={temp_act/real_temporal:.3f})"
              + (f"  [control's own temporal_activity={control_temp_act:.5f}]" if control_temp_act is not None else ""))

    (OUT_DIR / "comparison.json").write_text(json.dumps(results, indent=2))
    print(f"\nwrote {OUT_DIR / 'comparison.json'}")

    write_markdown_table(results)
    build_contact_sheet(real_matched, fg_mask)


def write_markdown_table(results):
    v = results["variants"]
    lines = []
    lines.append("# Wan2.1-VACE-1.3B control-format ablation -- metric comparison\n")
    if all(k in v for k in ("A", "B", "D")) and v["A"]["lpips"] and v["B"]["lpips"]:
        lpips_vace = [v[k]["lpips"]["overall"] for k in ("A", "A2", "D") if k in v and v[k]["lpips"]]
        lpips_ours = [v[k]["lpips"]["overall"] for k in ("B", "C") if k in v and v[k]["lpips"]]
        lines.append(
            f"## Verdict\n\n"
            f"**A vs B (load-bearing comparison, same model/scene/prompt/seed, only the control "
            f"representation differs): A looks and scores photoreal, B looks like a colourised "
            f"depth map and scores worse on every perceptual metric.** "
            f"Mean LPIPS with VACE's own depth format (A/A2/D) = {np.mean(lpips_vace):.3f}; "
            f"mean LPIPS with our depth format (B/C) = {np.mean(lpips_ours):.3f} "
            f"(higher = more different from real, worse). D (VACE's depth format + our plate.png "
            f"as reference image) is the best variant on every metric in this table. "
            f"This is our control-buffer format being out of distribution for VACE-1.3B, not the "
            f"model's ceiling -- the same checkpoint produces materially different, more photoreal "
            f"output when fed VACE's own depth representation instead of ours. "
            f"PSNR/SSIM alone would give the opposite (wrong) impression here -- B/C score "
            f"*higher* overall PSNR/SSIM than A/A2 despite looking like a false-colour depth-map "
            f"passthrough; LPIPS (perceptual) is what actually tracks the visual gap. See "
            f"`contact_sheet.png` for the visual evidence and the report for the full breakdown "
            f"(control-format finding) and the separate LEGO-styling / prompt-hijack finding "
            f"(A vs A2).\n"
        )
    lines.append(f"Ground truth: real DROID clip, {results['methodology']['n_frames_matched']} frames "
                  f"matched to each generated video via the upstream `VaceVideoProcessor` frame selection "
                  f"(see methodology in `compare_wan_variants.py` docstring and `frame_match.json`).\n")
    lines.append(f"LPIPS: {'alex net, spatial map, mask-averaged' if results['methodology']['lpips_available'] else 'NOT AVAILABLE, skipped (no substitute used)'}\n")
    lines.append(f"Real (matched-cadence) temporal activity baseline: "
                 f"**{results['real_baseline']['temporal_activity_matched_cadence']:.5f}** "
                 f"(mean abs frame-to-frame grayscale diff, [0,1] scale)\n")
    lines.append("| Variant | Control temporal activity | Output temporal activity | vs real | PSNR overall | PSNR bg | PSNR fg | SSIM overall | SSIM bg | SSIM fg | "
                  "LPIPS overall | LPIPS bg | LPIPS fg |")
    lines.append("|---|---|---|---|---|---|---|---|---|---|---|---|---|")
    for v, r in results["variants"].items():
        lp = r["lpips"]
        lp_o = f"{lp['overall']:.4f}" if lp else "n/a"
        lp_bg = f"{lp['background']:.4f}" if lp else "n/a"
        lp_fg = f"{lp['foreground']:.4f}" if lp else "n/a"
        ctrl_act = r.get("control_temporal_activity")
        ctrl_act_s = f"{ctrl_act:.5f}" if ctrl_act is not None else "n/a"
        lines.append(
            f"| {v} | {ctrl_act_s} | {r['temporal_activity']:.5f} | {r['temporal_activity_vs_real_ratio']:.3f}x "
            f"| {r['psnr_db']['overall']:.2f} | {r['psnr_db']['background']:.2f} | {r['psnr_db']['foreground']:.2f} "
            f"| {r['ssim']['overall']:.4f} | {r['ssim']['background']:.4f} | {r['ssim']['foreground']:.4f} "
            f"| {lp_o} | {lp_bg} | {lp_fg} |"
        )
    lines.append("\nNote: variants B and C's control video contains robot motion only (no brick/drawer -- "
                  "S6 object-pose wiring is a sibling's in-progress work), so their control's temporal "
                  "activity is expected to be well below A/D's dense-depth-network control and below the "
                  "real clip's, independent of any format quality question.")
    (OUT_DIR / "comparison.md").write_text("\n".join(lines) + "\n")
    print(f"wrote {OUT_DIR / 'comparison.md'}")


def build_contact_sheet(real_matched, fg_mask):
    variants_present = [v for v in VARIANTS if (OUT_DIR / v / "generated.mp4").exists()]
    rows = ["real"] + variants_present
    cols = CONTACT_SHEET_IDX
    oh, ow = real_matched.shape[1:3]
    pad = 4
    label_h = 24
    cell_w, cell_h = ow // 2, oh // 2  # downscale for a manageable contact sheet
    sheet_w = len(cols) * (cell_w + pad) + pad
    sheet_h = len(rows) * (cell_h + label_h + pad) + pad
    sheet = np.full((sheet_h, sheet_w, 3), 30, dtype=np.uint8)

    gen_cache = {}
    for v in variants_present:
        gen_cache[v] = np.stack(read_all_frames_rgb(OUT_DIR / v / "generated.mp4"))

    for ri, row in enumerate(rows):
        for ci, idx in enumerate(cols):
            if row == "real":
                frame = real_matched[idx]
            else:
                frames = gen_cache[row]
                frame = frames[idx] if idx < frames.shape[0] else frames[-1]
            thumb = cv2.resize(frame, (cell_w, cell_h))
            y0 = pad + ri * (cell_h + label_h + pad)
            x0 = pad + ci * (cell_w + pad)
            sheet[y0:y0 + cell_h, x0:x0 + cell_w] = thumb
            label = f"{row} @{idx}" if ci == 0 else f"@{idx}"
            cv2.putText(sheet, label, (x0, y0 + cell_h + 16), cv2.FONT_HERSHEY_SIMPLEX,
                        0.45, (255, 255, 255), 1, cv2.LINE_AA)
        # row label on the very left
        y0 = pad + ri * (cell_h + label_h + pad)
    out_path = OUT_DIR / "contact_sheet.png"
    cv2.imwrite(str(out_path), cv2.cvtColor(sheet, cv2.COLOR_RGB2BGR))
    print(f"wrote {out_path}")


if __name__ == "__main__":
    main()