Blame for scripts/teaser/compose_edit.py

#!/usr/bin/env python3
"""teaser · 世界编辑 case 的合成 / 打分 / 选片（maliang venv，CPU）。输入 = edits.py 的结果目录 <e81>/cases/<name>/（edit.json +
rollout.py 推完的 gen.mp4 / ids.npz / meta.json），原版 = edit.json 里 orig.roll（outputs/teaser/roll/<base>/gen.mp4，同首帧同种子）。

子命令
  side   一条 case：「原版 | 编辑版」并排（1920×1080），上方两张语义图（左 = 原版条件，右 = 编辑后条件，编辑区域高亮 + 白描边），
         编辑版上在编辑时刻弹出标签（如 "+ Tower @ 2.0 s"），底部时间条标出编辑帧；--hl N：编辑版视频上编辑区域描边 N 帧（默认 12）。
  grid   同一种编辑的若干 case 宫格（每格 = 编辑版 + 右上角编辑后语义小窗 + 底部标签；--pairs ⇒ 每格「原版 | 编辑版」两小格）。
  reel   若干 case 的 side 串接（相邻交叉淡化 --fade 帧；case 的 side 没有就先出）。
  sheet  一条 case 的帧拼图（原版 / 编辑版 / 编辑后语义 × 编辑前后 6 个时刻）→ <case>/edit_sheet.jpg（目测用）。
  score  编辑是否在视频里生效 → <case>/edit_score.json：
           * delta_in / delta_out：帧 ≥ 编辑结束（且 ≥ 编辑帧 + 4）的每 4 帧，编辑区域（掩码腐蚀 2 px）里编辑版与原版的 Lab ΔE
             均值 vs 区域外（掩码膨胀 24 px 之外、角色外）的 ΔE 均值；effect = delta_in − delta_out，ratio = in / out；
           * fit_edit / fit_orig：编辑区域里的 8×8 块按「本帧区域外各类的 Lab 外观（对角高斯）」分类，判成编辑后类 / 原类的比例
             （编辑后类在区域外不存在 ⇒ 记 None）；fit = fit_edit − fit_orig；
           * pre_delta：编辑帧之前的帧（≥ 4）整帧 ΔE（双向教师看全片条件 ⇒ 编辑前也会有点变化，因果学生应 ≈ 0）。
         有 vscore（outputs/teaser/vscore/edit/<name>.vscore.json）就一起汇总。
  index  全部 case → <out>/index.jsonl + index.md（编辑类型、标签、指标、路径），按 edit_score 排序。
  dmd    双向教师 vs 因果 DMD 学生 2×2 对照（<case>/edit_dmd_cmp.mp4；DMD 结果在 --dmd_root，默认 <e81>/dmd/<name>、<base>__orig）。

usage（$PY = .venv/bin/python，在 $M 下运行）：
  $PY scripts/teaser/compose_edit.py side outputs/teaser/edit/e81/cases/<name> [-o x.mp4]
  $PY scripts/teaser/compose_edit.py sheet outputs/teaser/edit/e81/cases/<name>
  $PY scripts/teaser/compose_edit.py score 'outputs/teaser/edit/e81/cases/*' [--procs 32]
  $PY scripts/teaser/compose_edit.py grid c1 c2 ... -o g.mp4 --cols 3 [--pairs]
  $PY scripts/teaser/compose_edit.py reel @list.txt -o r.mp4
  $PY scripts/teaser/compose_edit.py index 'outputs/teaser/edit/e81/cases/*' -o outputs/teaser/edit/e81/showcase
  $PY scripts/teaser/compose_edit.py dmd outputs/teaser/edit/e81/cases/<name> [--dmd_root DIR] [-o x.mp4]
glob 参数要加引号（脚本自己展开），也可以 @list.txt（每行一个 case 目录）。
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import sys
from concurrent.futures import ProcessPoolExecutor

import cv2
import numpy as np
from PIL import Image, ImageDraw

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
HERE = os.path.dirname(os.path.abspath(__file__))
for _p in (f"{M}/code/gc25d", HERE):
    if _p not in sys.path:
        sys.path.insert(0, _p)
import compose as CP  # noqa: E402
import geometry as GEO  # noqa: E402

FPS = 16
H, W = 480, 832
HL = (255, 214, 64)          # 编辑区域高亮色（暖黄）
ACCENT = (255, 196, 40)


def rp(p: str) -> str:
    return p if os.path.isabs(p) else os.path.join(M, p)


def expand(items):
    out = []
    for it in items:
        if it.startswith("@"):
            out += [ln.strip() for ln in open(it[1:]) if ln.strip() and not ln.startswith("#")]
        elif any(ch in it for ch in "*?["):
            out += sorted(glob.glob(rp(it)))
        else:
            out.append(it)
    return [rp(x) for x in out]


# ------------------------------------------------------------------ case ----
def load(cd: str, video: bool = True) -> dict:
    cd = rp(cd)
    e = json.load(open(os.path.join(cd, "edit.json")))
    z = np.load(os.path.join(cd, "mask.npz"))
    mask = np.unpackbits(z["bits"])[:int(np.prod(z["shape"]))].reshape(z["shape"]).astype(bool)
    c = dict(dir=cd, edit=e, mask=mask, ids1=np.load(os.path.join(cd, "ids.npz"))["ids"],
             ids0=np.load(rp(e["orig"]["ids"]))["ids"], name=e["name"])
    if video:
        c["gen1"] = CP.read_video(os.path.join(cd, "gen.mp4"))
        c["gen0"] = CP.read_video(rp(e["orig"]["gen"]))
    return c


def base_labels(e: dict) -> tuple[str, str]:
    b = e["base"]
    scene = b.get("scene") or ""
    try:
        scene = CP.SCENE_EN.get(b["scene"].split("@")[0], CP.pretty(b["scene"].split("@")[0]))
        look = CP.LOOK_EN.get(b["scene"].split("@")[1], "") if "@" in b["scene"] else ""
        scene = f"{scene} · {look}" if look else scene
    except Exception:                                                     # noqa: BLE001
        pass
    mp = CP.MAP_EN.get(b.get("theme") or "", CP.pretty(b.get("theme")))
    return scene, mp


_LF = {}


def label_pill(s: str, size: int, fg=(20, 20, 24), bg=ACCENT, alpha: float = 0.95):
    """编辑标签的小胶囊：拉丁 / 箭头走 DejaVu Sans Bold（Noto 变体字体没有 →），中文走 Droid。"""
    from PIL import ImageFont
    key = size
    if key not in _LF:
        _LF[key] = (ImageFont.truetype("/usr/share/fonts/dejavu-sans-fonts/DejaVuSans-Bold.ttf", size),
                    ImageFont.truetype(CP.FONT_CJK, size))
    lat, cjk = _LF[key]
    runs = CP._runs(s)
    w = sum((cjk if c else lat).getlength(r) for c, r in runs)
    px, py = round(size * 0.55), round(size * 0.3)
    W_, H_ = int(np.ceil(w)) + 2 * px, int(round(size * 1.1)) + 2 * py
    im = Image.new("RGBA", (W_, H_), (0, 0, 0, 0))
    d = ImageDraw.Draw(im)
    d.rounded_rectangle((0, 0, W_ - 1, H_ - 1), radius=H_ // 2, fill=(*bg, int(255 * alpha)))
    x = px
    for c, r in runs:
        f = cjk if c else lat
        d.text((x, py + 0.85 * size), r, font=f, fill=(*fg, 255), anchor="ls")
        x += f.getlength(r)
    a = np.asarray(im, np.float32)
    return a[..., :3].astype(np.uint8), a[..., 3:] / 255.0


def hl_sem(ids: np.ndarray, mask: np.ndarray, w: int, h: int, strength: float = 0.0) -> np.ndarray:
    """语义色图 + 编辑区域：区域外压暗 strength，区域描白边。"""
    im = CP.sem_panel(ids, w, h).astype(np.float32)
    m = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST) > 0
    if m.any():
        if strength > 0:
            im[~m] *= 1 - strength
        mm = cv2.dilate(m.astype(np.uint8), np.ones((3, 3), np.uint8))
        o = (cv2.dilate(mm, np.ones((3, 3), np.uint8)) - mm) > 0
        im[o] = (255, 255, 255)
    return im.clip(0, 255).astype(np.uint8)


def outline_on(img: np.ndarray, mask: np.ndarray, color=HL, alpha: float = 1.0, t: int = 2) -> np.ndarray:
    h, w = img.shape[:2]
    m = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
    cs, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cs = [c for c in cs if cv2.contourArea(c) > 30]
    if not cs:
        return img
    ov = img.copy()
    cv2.drawContours(ov, cs, -1, color, t, cv2.LINE_AA)
    return cv2.addWeighted(ov, alpha, img, 1 - alpha, 0)


def timeline(w: int, f: int, n: int, k: int, k1: int, h: int = 10) -> np.ndarray:
    bar = np.zeros((h, w, 3), np.uint8)
    bar[:] = (44, 48, 58)
    x = lambda i: int(round(i / max(1, n - 1) * (w - 1)))
    bar[:, :x(f) + 1] = (110, 118, 136)
    bar[:, x(k):x(max(k1, k)) + 2] = ACCENT
    return bar


# ------------------------------------------------------------------ side ----
def side_frames(c: dict, hl: int = 12, canvas=(1920, 1080)):
    e = c["edit"]
    Wc, Hc = canvas
    k, k1 = e["edit_frame"], max(e["edit_end"], e["edit_frame"])
    n = min(len(c["gen0"]), len(c["gen1"]), len(c["ids1"]))
    m, gap = 40, 40
    vw = (Wc - 2 * m - gap) // 2
    vh = int(round(vw * H / W))
    sw, sh = int(vw * 0.5), int(vw * 0.5 * H / W)
    top = 92
    y_sem = top
    y_vid = y_sem + sh + 46
    xs = [m, m + vw + gap]
    scene, mp = base_labels(e)
    head = Image.new("RGB", (Wc, top - 12), CP.BG)
    d = ImageDraw.Draw(head)
    x = CP.draw_text(d, m, 22, "Edit the map, the world follows", 30, CP.FG, "SemiBold")
    CP.draw_text(d, x + 24, 30, f"{scene}  ·  {mp} map", 20, CP.MUTED)
    head = np.asarray(head)
    lab = label_pill(e["label"], 26)
    lab_o = CP.pill("Original", 22, alpha=0.55)
    lab_e = CP.pill("Edited", 22, alpha=0.55)
    sem_l0 = CP.pill("original map", 17, alpha=0.6)
    sem_l1 = CP.pill("edited map", 17, alpha=0.6)
    for f in range(n):
        img = np.zeros((Hc, Wc, 3), np.uint8)
        img[:] = CP.BG
        img[12:12 + head.shape[0]] = head
        on = f >= k
        s0 = CP.sem_panel(c["ids0"][f], sw, sh)
        s1 = hl_sem(c["ids1"][f], c["mask"][f], sw, sh) if on else CP.sem_panel(c["ids1"][f], sw, sh)
        for i, (sp, sl) in enumerate(((s0, sem_l0), (s1, sem_l1))):
            xx = xs[i] + (vw - sw) // 2
            img[y_sem:y_sem + sh, xx:xx + sw] = sp
            CP.frame_rect(img, xx, y_sem, sw, sh)
            CP.paste_rgba(img, sl, xx + 8, y_sem + 8)
        g0 = CP.fit(c["gen0"][f], vw, vh)
        g1 = CP.fit(c["gen1"][f], vw, vh)
        if hl and k <= f < k1 + hl and c["mask"][f].any():
            a = 1.0 if f < k1 + hl // 2 else max(0.0, 1 - (f - k1 - hl // 2) / max(1, hl - hl // 2))
            g1 = outline_on(g1, c["mask"][f], HL, 0.85 * a, 2)
        for i, (g, lb) in enumerate(((g0, lab_o), (g1, lab_e))):
            img[y_vid:y_vid + vh, xs[i]:xs[i] + vw] = g
            CP.frame_rect(img, xs[i], y_vid, vw, vh, ACCENT if (i == 1 and on) else CP.BORDER, 2 if (i == 1 and on) else 1)
            CP.paste_rgba(img, lb, xs[i] + 12, y_vid + 12)
        if on:
            pop = min(1.0, (f - k + 1) / 4)
            rgb, al = lab
            CP.paste_rgba(img, (rgb, al * pop), xs[1] + vw - rgb.shape[1] - 14, y_vid + 12)
        bar = timeline(2 * vw + gap, f, n, k, k1)
        img[y_vid + vh + 18:y_vid + vh + 18 + bar.shape[0], m:m + bar.shape[1]] = bar
        rgb, al = CP.pill(f"{f / FPS:.2f} s", 18, alpha=0.0)
        CP.paste_rgba(img, (rgb, al), m, y_vid + vh + 32)
        yield img


def cmd_side(a) -> None:
    for cd in expand(a.cases):
        c = load(cd)
        out = a.o or os.path.join(cd, "edit_side.mp4")
        CP.write_video(out, side_frames(c, a.hl), crf=a.crf)
        print(out)


# ------------------------------------------------------------------ sheet ----
def sheet_idx(e: dict, n: int = 81):
    k, k1 = e["edit_frame"], max(e["edit_end"], e["edit_frame"])
    idx = [max(0, k - 6), min(n - 1, k + 3), min(n - 1, (k + k1) // 2 + 4), min(n - 1, k1 + 6), min(n - 1, (k1 + n) // 2 + 6), n - 1]
    return sorted(set(idx))


def cmd_sheet(a) -> None:
    for cd in expand(a.cases):
        c = load(cd)
        e = c["edit"]
        idx = sheet_idx(e, len(c["gen1"]))
        tw, th = W // 2, H // 2
        rows = []
        for kind in ("orig", "edit", "sem"):
            tiles = []
            for f in idx:
                if kind == "orig":
                    im = CP.fit(c["gen0"][f], tw, th)
                elif kind == "edit":
                    im = CP.fit(c["gen1"][f], tw, th)
                    if f >= e["edit_frame"] and c["mask"][f].any():
                        im = outline_on(im, c["mask"][f], HL, 0.9, 1)
                else:
                    im = hl_sem(c["ids1"][f], c["mask"][f], tw, th)
                im = im.copy()
                txt = f"{kind} f{f} {f / FPS:.2f}s"
                cv2.putText(im, txt, (6, th - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 3, cv2.LINE_AA)
                cv2.putText(im, txt, (6, th - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1, cv2.LINE_AA)
                tiles.append(im)
            rows.append(np.concatenate(tiles, 1))
        body = np.concatenate(rows, 0)
        head = Image.new("RGB", (body.shape[1], 28), CP.BG)
        sc = ""
        sp = os.path.join(cd, "edit_score.json")
        if os.path.exists(sp):
            s = json.load(open(sp))
            sc = f"  ·  effect {s.get('effect')}  ratio {s.get('ratio')}  fit {s.get('fit')}"
        CP.draw_text(ImageDraw.Draw(head), 8, 5, f"{e['name']}  ·  {e['label']}{sc}", 15, CP.FG, "SemiBold")
        out = a.o or os.path.join(cd, "edit_sheet.jpg")
        cv2.imwrite(out, cv2.cvtColor(np.concatenate([np.asarray(head), body], 0), cv2.COLOR_RGB2BGR),
                    [cv2.IMWRITE_JPEG_QUALITY, 88])
        print(out)


# ------------------------------------------------------------------ score ----
def _lab(rgb):
    return cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32) * np.array([100 / 255, 1, 1], np.float32) - [0, 128, 128]


def _collapse(ids):
    out = ids.copy()
    out[(ids >= 12) & (ids <= 16)] = 10
    out[(ids >= 17) & (ids <= 21)] = 11
    return out


def _fit_block(L, ids_edit, ids_orig, mask, cell=8):
    """编辑区域里的 8×8 块按区域外各类（编辑后语义）的 Lab 外观分类 → (判成编辑后类的比例, 判成原类的比例, 块数)。"""
    h, w = H // cell, W // cell
    blk = L[:h * cell, :w * cell].reshape(h, cell, w, cell, 3)
    feat = np.concatenate([blk.mean((1, 3)), blk[..., 0].std((1, 3))[..., None]], -1).reshape(-1, 4)

    def maj(ids):
        b = ids[:h * cell, :w * cell].reshape(h, cell, w, cell).transpose(0, 2, 1, 3).reshape(h * w, -1)
        mj = np.array([np.bincount(r, minlength=256).argmax() for r in b])
        return mj, (b == mj[:, None]).mean(1)
    me, pe = maj(ids_edit)
    mo, _ = maj(ids_orig)
    mb = mask[:h * cell, :w * cell].reshape(h, cell, w, cell).transpose(0, 2, 1, 3).reshape(h * w, -1).mean(1)
    inside = (mb >= 0.75) & (pe >= 0.75) & (me != mo)
    outside = (mb == 0) & (pe >= 0.75)
    if inside.sum() < 3:
        return None, None, int(inside.sum())
    cls = [c for c in np.unique(me[outside]) if (me[outside] == c).sum() >= 4]
    if int(np.bincount(me[inside]).argmax()) not in cls:
        return None, None, int(inside.sum())
    X = feat[outside & np.isin(me, cls)]
    y = me[outside & np.isin(me, cls)]
    mu = np.stack([X[y == c].mean(0) for c in cls])
    var = np.stack([X[y == c].var(0) for c in cls]) + 4.0
    Xi = feat[inside]
    ll = -0.5 * (((Xi[:, None] - mu[None]) ** 2 / var[None]).sum(-1) + np.log(var).sum(-1)[None])
    pred = np.asarray(cls)[ll.argmax(1)]
    return float((pred == me[inside]).mean()), float((pred == mo[inside]).mean()), int(inside.sum())


def score_one(cd: str) -> dict | None:
    cd = rp(cd)
    if not os.path.exists(os.path.join(cd, "gen.mp4")):
        return None
    c = load(cd)
    e = c["edit"]
    k, k1 = e["edit_frame"], max(e["edit_end"], e["edit_frame"])
    n = min(len(c["gen0"]), len(c["gen1"]))
    t0 = min(n - 1, max(k1, k + 4))
    frames = list(range(t0, n, 4)) or [n - 1]
    ins, outs, fe, fo, nb = [], [], [], [], []
    for f in frames:
        m = c["mask"][f]
        if m.sum() < 200:
            continue
        L0, L1 = _lab(c["gen0"][f]), _lab(c["gen1"][f])
        dE = np.linalg.norm(L1 - L0, axis=-1)
        mi = cv2.erode(m.astype(np.uint8), np.ones((5, 5), np.uint8)) > 0
        if mi.sum() < 100:
            mi = m
        chars = (_collapse(c["ids0"][f]) >= 10) | (_collapse(c["ids1"][f]) >= 10)
        mo = ~(cv2.dilate(m.astype(np.uint8), np.ones((49, 49), np.uint8)) > 0) & ~cv2.dilate(chars.astype(np.uint8),
                                                                                              np.ones((31, 31), np.uint8)).astype(bool)
        ins.append(float(dE[mi & ~chars].mean()) if (mi & ~chars).sum() > 50 else float(dE[mi].mean()))
        if mo.sum() > 2000:
            outs.append(float(dE[mo].mean()))
        a1, a0, b = _fit_block(L1, _collapse(c["ids1"][f]), _collapse(c["ids0"][f]), m)
        if a1 is not None:
            fe.append(a1)
            fo.append(a0)
            nb.append(b)
    pre = []
    for f in range(4, max(4, k - 1), 6):
        pre.append(float(np.linalg.norm(_lab(c["gen1"][f]) - _lab(c["gen0"][f]), axis=-1).mean()))
    d_in = float(np.mean(ins)) if ins else None
    d_out = float(np.mean(outs)) if outs else None
    r = dict(name=c["name"], edit=e["edit"], variant=e["variant"], label=e["label"], edit_frame=k, edit_end=k1,
             frames=frames, delta_in=None if d_in is None else round(d_in, 3), delta_out=None if d_out is None else round(d_out, 3),
             effect=None if d_in is None or d_out is None else round(d_in - d_out, 3),
             ratio=None if d_in is None or not d_out else round(d_in / d_out, 3),
             fit_edit=round(float(np.mean(fe)), 3) if fe else None, fit_orig=round(float(np.mean(fo)), 3) if fo else None,
             fit=round(float(np.mean(fe) - np.mean(fo)), 3) if fe else None, fit_blocks=int(np.mean(nb)) if nb else 0,
             pre_delta=round(float(np.mean(pre)), 3) if pre else None,
             mask_px_mean=round(float(c["mask"][t0:].sum((1, 2)).mean()), 1))
    vs = f"{M}/outputs/teaser/vscore/edit/{c['name']}.vscore.json"
    if os.path.exists(vs):
        v = json.load(open(vs))
        r.update(teaser=v.get("teaser"), flags=v.get("flags"))
    json.dump(r, open(os.path.join(cd, "edit_score.json"), "w"), indent=1)
    return r


def cmd_score(a) -> None:
    cds = expand(a.cases)
    with ProcessPoolExecutor(a.procs) as ex:
        for r in ex.map(score_one, cds):
            if r:
                print(json.dumps({k: r[k] for k in ("name", "effect", "ratio", "fit", "pre_delta")}))


# ------------------------------------------------------------------ grid / reel ----
def cmd_grid(a) -> None:
    cs = [load(cd) for cd in expand(a.cases)]
    Wc, Hc = a.canvas
    n = len(cs)
    cols = a.cols
    rows = (n + cols - 1) // cols
    cell_w = (Wc - 2 * a.margin - (cols - 1) * a.gap) // cols
    aspect = (2 * W + 8) / H if a.pairs else W / H
    cell_h = int(cell_w / aspect)
    if rows * cell_h + (rows - 1) * a.gap + 2 * a.margin + 60 > Hc:
        cell_h = (Hc - 2 * a.margin - 60 - (rows - 1) * a.gap) // rows
        cell_w = int(cell_h * aspect)
    x0 = (Wc - (cols * cell_w + (cols - 1) * a.gap)) // 2
    y0 = 60 + (Hc - 60 - (rows * cell_h + (rows - 1) * a.gap)) // 2
    labs = [label_pill(c["edit"]["label"], 16) for c in cs]
    title = Image.new("RGB", (Wc, 56), CP.BG)
    CP.draw_text(ImageDraw.Draw(title), a.margin, 14, a.title or "Edit the map, the world follows", 26, CP.FG, "SemiBold")
    title = np.asarray(title)
    nf = min(len(c["gen1"]) for c in cs)

    def frames():
        for f in range(nf):
            img = np.zeros((Hc, Wc, 3), np.uint8)
            img[:] = CP.BG
            img[:56] = title
            for i, c in enumerate(cs):
                r_, q = divmod(i, cols)
                x, y = x0 + q * (cell_w + a.gap), y0 + r_ * (cell_h + a.gap)
                e = c["edit"]
                on = f >= e["edit_frame"]
                if a.pairs:
                    hw = (cell_w - 4) // 2
                    img[y:y + cell_h, x:x + hw] = CP.fit(c["gen0"][f], hw, cell_h)
                    img[y:y + cell_h, x + hw + 4:x + hw + 4 + hw] = CP.fit(c["gen1"][f], hw, cell_h)
                    gx, gw = x + hw + 4, hw
                else:
                    img[y:y + cell_h, x:x + cell_w] = CP.fit(c["gen1"][f], cell_w, cell_h)
                    gx, gw = x, cell_w
                iw = int(gw * 0.3)
                ih = int(iw * H / W)
                ins = hl_sem(c["ids1"][f], c["mask"][f], iw, ih) if on else CP.sem_panel(c["ids1"][f], iw, ih)
                img[y + 6:y + 6 + ih, gx + gw - iw - 6:gx + gw - 6] = ins
                CP.frame_rect(img, gx + gw - iw - 6, y + 6, iw, ih, (255, 255, 255) if on else CP.BORDER, 1)
                CP.frame_rect(img, x, y, cell_w, cell_h, ACCENT if on else CP.BORDER, 1)
                if on:
                    rgb, al = labs[i]
                    CP.paste_rgba(img, (rgb, al * min(1.0, (f - e["edit_frame"] + 1) / 4)), gx + 8, y + cell_h - rgb.shape[0] - 8)
            yield img
    CP.write_video(a.o, frames(), crf=a.crf)
    print(a.o)


def cmd_reel(a) -> None:
    import subprocess
    vids = []
    for cd in expand(a.cases):
        if cd.endswith(".mp4"):
            vids.append(cd)
            continue
        sp = os.path.join(cd, "edit_side.mp4")
        if not os.path.exists(sp):
            CP.write_video(sp, side_frames(load(cd), a.hl), crf=16)
        vids.append(sp)
    cmd = [sys.executable, os.path.join(HERE, "compose.py"), "reel", *vids, "-o", a.o, "--fade", str(a.fade)]
    subprocess.run(cmd, check=True)
    print(a.o)


# ------------------------------------------------------------------ dmd ----
DMD_ROOT = f"{M}/outputs/teaser/edit/e81/dmd"


def cmd_dmd(a) -> None:
    """双向教师 vs 因果 DMD 学生（同首帧、同条件、同种子）：2×2 =「双向 原版 | 双向 编辑版」/「DMD 原版 | DMD 编辑版」+ 编辑后语义小窗。
    DMD 结果 = <dmd_root>/<name>/gen.mp4 与 <dmd_root>/<base>__orig/gen.mp4（scripts/teaser/rollout_dmd.py 出的）。"""
    for cd in expand(a.cases):
        c = load(cd)
        e = c["edit"]
        root = a.dmd_root or DMD_ROOT
        dp, do = os.path.join(root, c["name"], "gen.mp4"), os.path.join(root, e["base"]["name"] + "__orig", "gen.mp4")
        if not (os.path.exists(dp) and os.path.exists(do)):
            print(f"[skip] 没有 DMD 结果：{c['name']}")
            continue
        d1, d0 = CP.read_video(dp), CP.read_video(do)
        n = min(len(d1), len(d0), len(c["gen1"]), len(c["gen0"]))
        Wc, Hc = 1920, 1080
        vw, vh = 780, 450
        xs, ys = (160, 980), (92, 92 + vh + 30)
        k, k1 = e["edit_frame"], max(e["edit_end"], e["edit_frame"])
        head = Image.new("RGB", (Wc, 80), CP.BG)
        dr = ImageDraw.Draw(head)
        x = CP.draw_text(dr, 160, 22, "Bidirectional teacher vs causal student", 28, CP.FG, "SemiBold")
        CP.draw_text(dr, x + 24, 29, e["label"], 20, CP.MUTED)
        head = np.asarray(head)
        tags = [[CP.pill(t, 20, alpha=0.6) for t in row] for row in (("Teacher · original", "Teacher · edited"),
                                                                       ("Student (causal) · original", "Student (causal) · edited"))]
        lab = label_pill(e["label"], 22)

        def frames():
            for f in range(n):
                img = np.zeros((Hc, Wc, 3), np.uint8)
                img[:] = CP.BG
                img[:80] = head
                for r, (g0, g1) in enumerate(((c["gen0"], c["gen1"]), (d0, d1))):
                    for q, g in enumerate((g0, g1)):
                        im = CP.fit(g[f], vw, vh)
                        if q == 1 and k <= f < k1 + 12 and c["mask"][f].any():
                            im = outline_on(im, c["mask"][f], HL, 0.8, 2)
                        img[ys[r]:ys[r] + vh, xs[q]:xs[q] + vw] = im
                        on = q == 1 and f >= k
                        CP.frame_rect(img, xs[q], ys[r], vw, vh, ACCENT if on else CP.BORDER, 2 if on else 1)
                        CP.paste_rgba(img, tags[r][q], xs[q] + 10, ys[r] + 10)
                        if on:
                            rgb, al = lab
                            CP.paste_rgba(img, (rgb, al * min(1.0, (f - k + 1) / 4)), xs[q] + vw - rgb.shape[1] - 10, ys[r] + 10)
                iw, ih = 224, 129
                ins = hl_sem(c["ids1"][f], c["mask"][f], iw, ih) if f >= k else CP.sem_panel(c["ids1"][f], iw, ih)
                img[ys[1] + vh - ih - 10:ys[1] + vh - 10, xs[1] + vw - iw - 10:xs[1] + vw - 10] = ins
                bar = timeline(xs[1] + vw - xs[0], f, n, k, k1, 8)
                img[ys[1] + vh + 14:ys[1] + vh + 22, xs[0]:xs[0] + bar.shape[1]] = bar
                yield img
        out = a.o or os.path.join(cd, "edit_dmd_cmp.mp4")
        CP.write_video(out, frames(), crf=a.crf)
        print(out)


# ------------------------------------------------------------------ index ----
def cmd_index(a) -> None:
    rows = []
    for cd in expand(a.cases):
        ep = os.path.join(cd, "edit.json")
        if not os.path.exists(ep):
            continue
        e = json.load(open(ep))
        r = dict(name=e["name"], dir=cd, edit=e["edit"], variant=e["variant"], label=e["label"], k=e["edit_frame"],
                 k_end=e["edit_end"], base=e["base"]["name"], theme=e["base"]["theme"], scene=e["base"]["scene"],
                 rolled=os.path.exists(os.path.join(cd, "meta.json")), weak=e.get("weak"),
                 orig=e["orig"]["gen"], gen=os.path.join(cd, "gen.mp4"))
        sp = os.path.join(cd, "edit_score.json")
        if os.path.exists(sp):
            s = json.load(open(sp))
            r.update({k: s.get(k) for k in ("effect", "ratio", "fit", "fit_edit", "fit_orig", "pre_delta", "teaser", "flags")})
        rows.append(r)
    os.makedirs(a.o, exist_ok=True)
    rows.sort(key=lambda r: (r["edit"], -(r.get("effect") or -99)))
    with open(os.path.join(a.o, "index.jsonl"), "w") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    with open(os.path.join(a.o, "index.md"), "w") as fh:
        fh.write(f"# 世界编辑 case（{len(rows)} 条）\n\n| edit | 标签 | effect | ratio | fit | pre | teaser | case |\n|---|---|---|---|---|---|---|---|\n")
        for r in rows:
            fh.write(f"| {r['edit']} | {r['label']} | {r.get('effect')} | {r.get('ratio')} | {r.get('fit')} | {r.get('pre_delta')} | "
                     f"{r.get('teaser')} | `{r['name']}` |\n")
    print(f"index：{len(rows)} 条 → {a.o}/index.md")


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(description="teaser 世界编辑合成 / 打分")
    ap.add_argument("cmd", choices=["side", "grid", "reel", "sheet", "score", "index", "dmd"])
    ap.add_argument("cases", nargs="+")
    ap.add_argument("-o", default=None)
    ap.add_argument("--hl", type=int, default=12)
    ap.add_argument("--crf", type=int, default=16)
    ap.add_argument("--cols", type=int, default=3)
    ap.add_argument("--pairs", action="store_true")
    ap.add_argument("--margin", type=int, default=32)
    ap.add_argument("--gap", type=int, default=16)
    ap.add_argument("--canvas", type=lambda s: tuple(int(x) for x in s.split("x")), default=(1920, 1080))
    ap.add_argument("--title", default=None)
    ap.add_argument("--fade", type=int, default=8)
    ap.add_argument("--procs", type=int, default=24)
    ap.add_argument("--dmd_root", default=None)
    return ap


def main(argv=None) -> None:
    a = build_parser().parse_args(argv)
    dict(side=cmd_side, grid=cmd_grid, reel=cmd_reel, sheet=cmd_sheet, score=cmd_score, index=cmd_index, dmd=cmd_dmd)[a.cmd](a)


if __name__ == "__main__":
    main()