Blame for scripts/genbench2_compose.py

#!/usr/bin/env python3
"""genbench2 结果整理：
* compare2_<tag>.mp4：每个地图一个 4 列网格（每格 416×240 半分辨率）：输入 ids / GT（val）/ 原首帧生成 / 12 种风格；
* compare2_L3.mp4：4 张新地图，每张一行块：ids / 第一轮原首帧生成 / 6 种风格；
* style_sheet.png：每个地图两行（帧 0 = FLUX 首帧；帧 80 = 生成的最后一帧），看风格能不能撑到最后；
* style_stats.json：风格保持 / 回退统计（HSV 直方图交）：
    keep  = sim(生成帧 t, 风格首帧)        越高越好（风格保持）
    back  = sim(生成帧 t, 原始未风格化帧 0) 相对 keep 越高 ⇒ 越像回退到训练外观
    以及 FLUX restyle 的结构保持分（firstframe_flux 的 struct_metrics）。"""
import json
import os
import sys

import cv2
import imageio.v3 as iio
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import genbench2_run as R  # noqa: E402

G, G1 = R.G, R.G1
L1 = f"{R.M}/outputs/sft_v10/eval_step-8000/L1_cond"
W, H = 832, 480
HW, HH = W // 2, H // 2


def panel(path, k):
    v = iio.imread(path)
    return v[:, -H:, k * W:(k + 1) * W]


def gen_frames(out_dir, name):
    p = f"{out_dir}/{name}_gen.mp4"
    return iio.imread(p) if os.path.exists(p) else panel(f"{out_dir}/{name}.mp4", 0)


def half(fr):
    return cv2.resize(fr, (HW, HH), interpolation=cv2.INTER_AREA)


def label(fr, text, scale=0.5):
    out = fr.copy()
    cv2.rectangle(out, (0, 0), (out.shape[1], 22), (0, 0, 0), -1)
    cv2.putText(out, text, (6, 16), cv2.FONT_HERSHEY_SIMPLEX, scale, (255, 230, 120), 1, cv2.LINE_AA)
    return out


def grid_video(tiles, ncol=4):
    T = min(len(v) for _, v in tiles)
    while len(tiles) % ncol:
        tiles.append(("", np.zeros((T, HH, HW, 3), np.uint8)))
    frames = []
    for t in range(T):
        rows = [np.concatenate([label(half(v[t]) if v.shape[1] != HH else v[t], n) for n, v in tiles[r:r + ncol]], 1)
                for r in range(0, len(tiles), ncol)]
        frames.append(np.concatenate(rows, 0))
    return np.stack(frames)


def write(path, frames):
    iio.imwrite(path, frames, fps=16, codec="libx264", quality=None, macro_block_size=16,
                ffmpeg_params=["-crf", "22", "-pix_fmt", "yuv420p", "-movflags", "+faststart"])
    print(path, frames.shape, round(os.path.getsize(path) / 1e6, 1), "MB", flush=True)


def hist(fr):
    hsv = cv2.cvtColor(fr, cv2.COLOR_RGB2HSV)
    h = cv2.calcHist([hsv], [0, 1, 2], None, [12, 6, 6], [0, 180, 0, 256, 0, 256]).ravel()
    return h / max(h.sum(), 1)


def sim(a, b):
    return float(np.minimum(hist(a), hist(b)).sum())


def main():
    stats, sheet_rows = {}, []
    TW, TH = 208, 120
    for tag in R.SOURCES:
        d, sd = f"{G}/L2/{tag}", R.src_dir(tag)
        orig0 = cv2.cvtColor(cv2.imread(R.pick(sd, "frame0.png")), cv2.COLOR_BGR2RGB)
        tiles = [("input ids (numeric, parts coloured)", panel(f"{d}/{tag}_{next(iter(R.STYLES))}.mp4", 1))]
        if tag.startswith("val"):
            cid = R.SOURCES[tag]
            tiles.append(("GT (held-out run)", panel(f"{L1}/{cid}.mp4", 0)))
            tiles.append(("gen: original first frame", panel(f"{L1}/{cid}.mp4", 1)))
        else:
            od = f"{G1}/L2/{tag}" if os.path.exists(f"{G1}/L2/{tag}/{tag}_orig.mp4") else d
            tiles.append(("gen: original first frame", gen_frames(od, f"{tag}_orig")))
        row0, row80 = [cv2.resize(orig0, (TW, TH), interpolation=cv2.INTER_AREA)], [None]
        stats[tag] = {}
        for s in R.STYLES:
            if not os.path.exists(f"{d}/{tag}_{s}.mp4"):
                print("缺", tag, s)
                continue
            v = gen_frames(d, f"{tag}_{s}")
            tiles.append((f"UNSEEN style: {s}", v))
            f0 = cv2.cvtColor(cv2.imread(f"{d}/{s}_f0.png"), cv2.COLOR_BGR2RGB)
            fj = json.load(open(f"{d}/{s}_f0.json"))
            stats[tag][s] = {"keep_t40": round(sim(v[40], f0), 3), "keep_t80": round(sim(v[80], f0), 3),
                             "back_t80": round(sim(v[80], orig0), 3), "flux_struct": fj.get("struct")}
            row0.append(cv2.resize(f0, (TW, TH), interpolation=cv2.INTER_AREA))
            row80.append(cv2.resize(v[80], (TW, TH), interpolation=cv2.INTER_AREA))
        write(f"{G}/compare2_{tag}.mp4", grid_video(tiles))
        o = tiles[2 if tag.startswith("val") else 1][1]
        row80[0] = cv2.resize(o[80], (TW, TH), interpolation=cv2.INTER_AREA)
        stats[tag]["_orig"] = {"keep_t80": round(sim(o[80], orig0), 3)}
        sheet_rows.append(label(np.concatenate(row0, 1), f"{tag}  frame 0 (FLUX first frames; col 1 = original)", 0.45))
        sheet_rows.append(label(np.concatenate(row80, 1), f"{tag}  frame 80 (generated)", 0.45))

    # L3
    blocks = []
    for c in sorted(os.listdir(f"{G}/L3")):
        d = f"{G}/L3/{c}"
        theme = c.split("__")[0]
        tiles = [(f"UNSEEN map [{theme}]: ids", panel(f"{G1}/L3/{c}/{c}.mp4", 1)),
                 ("gen: FLUX map2frame first frame", panel(f"{G1}/L3/{c}/{c}.mp4", 0))]
        orig0 = cv2.cvtColor(cv2.imread(f"{G1}/L3/{c}/f0.png"), cv2.COLOR_BGR2RGB)
        stats[f"L3_{theme}"] = {}
        for s in R.L3_STYLES:
            if os.path.exists(f"{d}/{c}_{s}.mp4"):
                v = gen_frames(d, f"{c}_{s}")
                tiles.append((f"UNSEEN style: {s}", v))
                f0 = cv2.cvtColor(cv2.imread(f"{d}/{s}_f0.png"), cv2.COLOR_BGR2RGB)
                stats[f"L3_{theme}"][s] = {"keep_t80": round(sim(v[80], f0), 3), "back_t80": round(sim(v[80], orig0), 3)}
        blocks.append(grid_video(tiles))
    if blocks:
        T = min(len(b) for b in blocks)
        write(f"{G}/compare2_L3.mp4", np.concatenate([b[:T] for b in blocks], 1))

    w = max(r.shape[1] for r in sheet_rows)
    sheet = np.concatenate([cv2.copyMakeBorder(r, 0, 4 * (i % 2), 0, w - r.shape[1], cv2.BORDER_CONSTANT)
                            for i, r in enumerate(sheet_rows)], 0)
    cv2.imwrite(f"{G}/style_sheet.jpg", sheet[:, :, ::-1], [cv2.IMWRITE_JPEG_QUALITY, 92])
    print("sheet", sheet.shape)

    # 汇总：每种风格跨地图平均
    agg = {}
    for tag, d in stats.items():
        for s, m in d.items():
            if s.startswith("_"):
                continue
            a = agg.setdefault(s, {"keep_t80": [], "back_t80": []})
            a["keep_t80"].append(m["keep_t80"]); a["back_t80"].append(m["back_t80"])
    summary = {s: {k: round(float(np.mean(v)), 3) for k, v in a.items()} | {"n": len(a["keep_t80"])} for s, a in agg.items()}
    summary["_orig_keep_t80"] = round(float(np.mean([d["_orig"]["keep_t80"] for d in stats.values() if "_orig" in d])), 3)
    json.dump({"per_style": summary, "per_run": stats}, open(f"{G}/style_stats.json", "w"), indent=1)
    print(json.dumps(summary, indent=1))


if __name__ == "__main__":
    sys.exit(main())