processing/build_camera_dataset.py

#!/usr/bin/env python3
"""
Build LTX-2/Camera-Dataset from the Kwai CameraClone-Dataset.

Produces reference/target video pairs that share a camera trajectory but show
different scenes, in the flat `<name>.mp4` / `<name>_camera.mp4` + dataset.json
layout used by LTX-2/Depth-Dataset, Pose-Dataset and Inpainting-Dataset.

Pairing comes from the source dataset's own structure: within one `traj_<group>_<cam>`
directory every video was rendered with the same camera trajectory at a different
3D location, so (video_path, ref_video_path) from a CSV row is a valid pair by
construction. See README.md for the measurements behind the design.

Stages (all resumable; artifacts cached in this directory):
    index    parse the 1.3 GB CSV into a compact path index
    select   caption-cluster the candidates and pick 2400 scene-diverse pairs
    measure  estimate camera motion for the selected targets
    trim     cap near-static trajectories, cut to the final 2000
    encode   normalise both videos of every pair to 1280x720 / 73 frames / h264
    emit     write dataset.json and the file -> original path mapping CSV
    verify   check counts, formats, pair correctness and diversity
"""

import argparse
import collections
import csv
import json
import os
import random
import re
import subprocess
import sys
import time

import numpy as np
import pandas as pd

SRC_ROOT = "/mnt/dataset/xinyuy/datasets/CameraClone-Dataset"
SRC_CSV = f"{SRC_ROOT}/CamCloneDataset.csv"
WORK = f"{SRC_ROOT}/processing"
LTX = "/mnt/dataset/xinyuy/datasets/LTX-2"

# The val split excludes every trajectory group and every scene location used by train,
# so the two sets share no camera trajectory and no 3D location — zero overlap.
SPLITS = {
    "train": dict(out_dir=f"{LTX}/Camera-Dataset", sfx="",
                  n_final=2000, n_candidates=2400, exclude=None),
    "val":   dict(out_dir=f"{LTX}/Camera-Dataset-Val", sfx="_val",
                  n_final=50, n_candidates=70, exclude="train"),
}

STATIC_CAP = 0.08            # max share of near-static trajectories in the final set
N_CLUSTERS = 60              # caption clusters -> scene diversity
SEED = 0

# Native mode remuxes the source video stream untouched: HEVC / yuvj420p / 15 fps /
# 77 frames at the batch's own resolution. Set NATIVE = False to re-encode to
# OUT_W x OUT_H at OUT_FRAMES (h264) instead.
NATIVE = True
NATIVE_FRAMES = 77
NATIVE_CODEC, NATIVE_PIX_FMT = "hevc", "yuvj420p"
NATIVE_DIMS = {"0316": (1344, 768), "0317": (1344, 768)}
NATIVE_DIMS_DEFAULT = (1008, 576)

OUT_W, OUT_H, OUT_FRAMES = 1280, 720, 73   # 73 = 8*9+1, LTX latent-safe under the native 77
ENCODE_JOBS, ENCODE_THREADS = 32, 3
MEASURE_WORKERS = 64

# motion thresholds, in units of the 320px-wide analysis frame
T_PAN, T_TILT, T_ROLL, T_ZOOM = 30.0, 30.0, 8.0, 15.0
T_LINEAR, T_ARC = 0.85, 0.40

PATH_RE = re.compile(r"^\./data/(\d+)/traj_(\d+)_(\d+)/scene(\d+)_(\d+)\.mp4$")

# index and candidates (with captions + clusters) are split-independent and shared.
INDEX_PQ = f"{WORK}/index.parquet"
CAND_PQ = f"{WORK}/candidates.parquet"

# Per-split; bound by configure().
SPLIT = OUT_DIR = SEL_CSV = MOTION_CSV = FINAL_CSV = MAPPING_CSV = None
N_FINAL = N_CANDIDATES = EXCLUDE_CSV = None


def configure(split):
    """Bind the per-split output paths and sizes."""
    global SPLIT, OUT_DIR, SEL_CSV, MOTION_CSV, FINAL_CSV, MAPPING_CSV
    global N_FINAL, N_CANDIDATES, EXCLUDE_CSV
    cfg = SPLITS[split]
    SPLIT, OUT_DIR = split, cfg["out_dir"]
    N_FINAL, N_CANDIDATES = cfg["n_final"], cfg["n_candidates"]
    sfx = cfg["sfx"]
    SEL_CSV = f"{WORK}/selected_candidates{sfx}.csv"
    MOTION_CSV = f"{WORK}/motion_labels{sfx}.csv"
    FINAL_CSV = f"{WORK}/final_pairs{sfx}.csv"
    MAPPING_CSV = f"{WORK}/camera_dataset_mapping{sfx}.csv"
    EXCLUDE_CSV = (f"{WORK}/final_pairs{SPLITS[cfg['exclude']]['sfx']}.csv"
                   if cfg["exclude"] else None)


def log(msg):
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


def native_dims(date):
    """Source resolution of a render batch. Both videos of a pair come from the same
    trajectory dir, hence the same batch, so a pair is always internally consistent."""
    return NATIVE_DIMS.get(str(date), NATIVE_DIMS_DEFAULT)


# --------------------------------------------------------------------------- index

def stage_index():
    """Stream the source CSV into a compact index. Captions are left on disk."""
    log(f"indexing {SRC_CSV}")
    rows = []
    bad = 0
    with open(SRC_CSV) as f:
        f.readline()  # header
        for i, line in enumerate(f):
            # video_path and ref_video_path are unquoted and contain no commas;
            # the file has no embedded newlines (verified: lines == rows + 1).
            a, b, _ = line.split(",", 2)
            ma, mb = PATH_RE.match(a), PATH_RE.match(b)
            if not ma or not mb:
                bad += 1
                continue
            date, grp, cam, tsid, _ = ma.groups()
            rows.append((i, date, int(grp), int(cam), int(tsid), int(mb.group(4))))
    df = pd.DataFrame(rows, columns=["row", "date", "grp", "cam", "tsid", "rsid"])
    df.to_parquet(INDEX_PQ, index=False)
    log(f"indexed {len(df):,} rows ({bad} unparseable) -> {INDEX_PQ}")
    log(f"  {df.groupby(['date', 'grp']).ngroups:,} trajectory groups, "
        f"{df.groupby(['date', 'tsid']).ngroups:,} target locations")


# --------------------------------------------------------------------------- select

def _fetch_captions(rownos):
    """Second pass over the CSV, pulling captions only for the rows we need."""
    want = set(rownos)
    out = {}
    csv.field_size_limit(10 ** 7)
    with open(SRC_CSV, newline="") as f:
        for i, rec in enumerate(csv.reader(f)):
            if i == 0:
                cap_col = rec.index("caption")
                continue
            j = i - 1
            if j in want:
                out[j] = rec[cap_col]
                if len(out) == len(want):
                    break
    return out


def stage_select():
    rng = random.Random(SEED)
    df = pd.read_parquet(INDEX_PQ)
    log(f"loaded index: {len(df):,} rows")

    if os.path.exists(CAND_PQ):
        cand = pd.read_parquet(CAND_PQ)
        log(f"reusing cached candidate pool: {len(cand):,} candidates with captions + clusters")
    else:
        # One candidate per (group, target location): each group's 4 locations are
        # different scenes, so this gives up to 4 caption choices per trajectory group.
        log("building candidate pool (one row per group x target location)")
        cand = (df.sample(frac=1.0, random_state=SEED)
                  .drop_duplicates(subset=["date", "grp", "tsid"], keep="first")
                  .reset_index(drop=True))
        log(f"  {len(cand):,} candidates over {cand.groupby(['date', 'grp']).ngroups:,} groups")

        log("fetching captions for candidates (second CSV pass)")
        caps = _fetch_captions(cand["row"].tolist())
        cand["caption"] = cand["row"].map(caps)
        missing = cand["caption"].isna().sum()
        if missing:
            log(f"  dropping {missing} candidates with no caption")
            cand = cand[cand["caption"].notna()].reset_index(drop=True)

        log(f"clustering captions into {N_CLUSTERS} scene clusters (TF-IDF + MiniBatchKMeans)")
        from sklearn.feature_extraction.text import TfidfVectorizer
        from sklearn.cluster import MiniBatchKMeans
        vec = TfidfVectorizer(stop_words="english", ngram_range=(1, 2),
                              max_features=50000, min_df=5, max_df=0.5)
        X = vec.fit_transform(cand["caption"])
        km = MiniBatchKMeans(n_clusters=N_CLUSTERS, random_state=SEED,
                             n_init=5, batch_size=4096)
        cand["cluster"] = km.fit_predict(X)
        cand.to_parquet(CAND_PQ, index=False)
        log(f"  clusters sized {cand['cluster'].value_counts().min()}"
            f"..{cand['cluster'].value_counts().max()}")

    # Proportional per-date quota, so no render batch dominates.
    groups_per_date = df.groupby("date")["grp"].nunique()
    total_groups = groups_per_date.sum()
    quota = {d: int(round(N_CANDIDATES * n / total_groups))
             for d, n in groups_per_date.items()}
    log(f"  per-date quota: {quota}")

    # Round-robin over caption clusters, always drawing from the least-used one.
    pools = collections.defaultdict(list)
    for rec in cand.sample(frac=1.0, random_state=SEED).itertuples(index=False):
        pools[rec.cluster].append(rec)
    cursor = collections.defaultdict(int)
    taken = collections.Counter()
    used_groups, used_tloc, used_rloc = set(), set(), set()
    date_count = collections.Counter()
    selected = []

    # Pre-seeding the "used" sets with the other split's picks is what makes the splits
    # disjoint: no shared trajectory group, no shared target or reference location.
    if EXCLUDE_CSV:
        if not os.path.exists(EXCLUDE_CSV):
            sys.exit(f"split '{SPLIT}' excludes {EXCLUDE_CSV}, which does not exist yet")
        prev = pd.read_csv(EXCLUDE_CSV, dtype={"date": str})
        for r in prev.itertuples(index=False):
            used_groups.add((r.date, r.grp))
            used_tloc.add((r.date, r.tsid))
            used_rloc.add((r.date, r.rsid))
        log(f"  holding out {len(used_groups):,} trajectory groups and "
            f"{len(used_tloc | used_rloc):,} locations used by {os.path.basename(EXCLUDE_CSV)}")

    while len(selected) < N_CANDIDATES:
        progressed = False
        for cl in sorted(pools, key=lambda c: (taken[c], rng.random())):
            pool, i = pools[cl], cursor[cl]
            while i < len(pool):
                r = pool[i]
                i += 1
                if ((r.date, r.grp) in used_groups
                        or (r.date, r.tsid) in used_tloc
                        or (r.date, r.rsid) in used_rloc
                        or date_count[r.date] >= quota.get(r.date, 0)):
                    continue
                used_groups.add((r.date, r.grp))
                used_tloc.add((r.date, r.tsid))
                used_rloc.add((r.date, r.rsid))
                date_count[r.date] += 1
                taken[cl] += 1
                selected.append(r)
                progressed = True
                break
            cursor[cl] = i
            if progressed:
                break
        if not progressed:
            log(f"  pool exhausted at {len(selected)} candidates")
            break

    sel = pd.DataFrame(selected)
    sel["target_src"] = sel.apply(
        lambda r: f"data/{r.date}/traj_{r.grp}_{r.cam:02d}/scene{r.tsid}_{r.cam:02d}.mp4", axis=1)
    sel["ref_src"] = sel.apply(
        lambda r: f"data/{r.date}/traj_{r.grp}_{r.cam:02d}/scene{r.rsid}_{r.cam:02d}.mp4", axis=1)
    sel["name"] = sel.apply(
        lambda r: f"{r.date}_g{r.grp}_c{r.cam:02d}_s{r.tsid}", axis=1)
    sel.to_csv(SEL_CSV, index=False)
    log(f"selected {len(sel):,} candidates -> {SEL_CSV}")
    log(f"  dates: {dict(sel['date'].value_counts().sort_index())}")
    log(f"  clusters used: {sel['cluster'].nunique()}/{N_CLUSTERS}, "
        f"max per cluster {sel['cluster'].value_counts().max()}")


# --------------------------------------------------------------------------- measure

def measure_trajectory(path, step=4, width=320):
    """Cumulative 2D camera motion via ORB + RANSAC similarity fit.

    RANSAC rejects the animated character, so the fit tracks the background —
    i.e. the camera. Returns [pan, tilt, roll_deg, 100*log_scale, straightness],
    where straightness = |net displacement| / path length separates linear moves
    (~1.0) from arcs and splines.
    """
    import cv2
    cv2.setNumThreads(1)
    cap = cv2.VideoCapture(path)
    frames, i = [], 0
    while True:
        ok, fr = cap.read()
        if not ok:
            break
        if i % step == 0:
            h0, w0 = fr.shape[:2]
            fr = cv2.resize(fr, (width, int(width * h0 / w0)))
            frames.append(cv2.cvtColor(fr, cv2.COLOR_BGR2GRAY))
        i += 1
    cap.release()
    if len(frames) < 3:
        return None

    orb = cv2.ORB_create(1500)
    kd = [orb.detectAndCompute(f, None) for f in frames]
    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    pan = tilt = roll = lsc = 0.0
    path_len = 0.0
    n = 0
    for (k1, d1), (k2, d2) in zip(kd[:-1], kd[1:]):
        if d1 is None or d2 is None or len(k1) < 8 or len(k2) < 8:
            continue
        m = bf.match(d1, d2)
        if len(m) < 8:
            continue
        src = np.float32([k1[x.queryIdx].pt for x in m]).reshape(-1, 1, 2)
        dst = np.float32([k2[x.trainIdx].pt for x in m]).reshape(-1, 1, 2)
        M, _ = cv2.estimateAffinePartial2D(src, dst, method=cv2.RANSAC,
                                           ransacReprojThreshold=2.0)
        if M is None:
            continue
        a, b = M[0, 0], M[1, 0]
        dx, dy = M[0, 2], M[1, 2]
        pan += dx
        tilt += dy
        roll += np.degrees(np.arctan2(b, a))
        lsc += np.log(max(np.hypot(a, b), 1e-6))
        path_len += np.hypot(dx, dy)
        n += 1
    if n == 0:
        return None
    net = np.hypot(pan, tilt)
    straight = net / path_len if path_len > 1e-6 else 0.0
    return [pan, tilt, roll, lsc * 100, straight]


def _measure_one(args):
    name, path = args
    try:
        return name, measure_trajectory(path)
    except Exception:
        return name, None


def classify(pan, tilt, roll, zoom, straight):
    comps = []
    if abs(pan) > T_PAN:
        comps.append("pan_right" if pan < 0 else "pan_left")
    if abs(tilt) > T_TILT:
        comps.append("tilt_down" if tilt < 0 else "tilt_up")
    if abs(roll) > T_ROLL:
        comps.append("roll_cw" if roll < 0 else "roll_ccw")
    if abs(zoom) > T_ZOOM:
        comps.append("zoom_out" if zoom < 0 else "zoom_in")
    n_axes = len(comps)
    if not comps:
        comps = ["static"]
    shape = "linear" if straight >= T_LINEAR else ("arc" if straight >= T_ARC else "complex")
    return "+".join(sorted(comps)), n_axes, shape


def stage_measure():
    from concurrent.futures import ProcessPoolExecutor
    sel = pd.read_csv(SEL_CSV, dtype={"date": str})
    done = {}
    if os.path.exists(MOTION_CSV):
        prev = pd.read_csv(MOTION_CSV)
        done = set(prev["name"])
        log(f"resuming: {len(done)} already measured")
    todo = [(r.name_, f"{SRC_ROOT}/{r.target_src}")
            for r in sel.rename(columns={"name": "name_"}).itertuples(index=False)
            if r.name_ not in done]
    log(f"measuring camera motion for {len(todo):,} targets on {MEASURE_WORKERS} workers")

    recs = []
    t0 = time.time()
    with ProcessPoolExecutor(MEASURE_WORKERS) as ex:
        for k, (name, sig) in enumerate(ex.map(_measure_one, todo, chunksize=4), 1):
            if sig is not None:
                comps, n_axes, shape = classify(*sig)
                recs.append(dict(name=name, pan=sig[0], tilt=sig[1], roll=sig[2],
                                 zoom=sig[3], straightness=sig[4],
                                 motion_components=comps, n_axes=n_axes, path_shape=shape))
            if k % 400 == 0:
                log(f"  {k}/{len(todo)} ({time.time() - t0:.0f}s)")
    new = pd.DataFrame(recs)
    if done:
        new = pd.concat([prev, new], ignore_index=True)
    new.to_csv(MOTION_CSV, index=False)
    log(f"measured {len(new):,} trajectories in {time.time() - t0:.0f}s -> {MOTION_CSV}")
    log(f"  static share: {(new['motion_components'] == 'static').mean():.1%}")


# --------------------------------------------------------------------------- trim

def stage_trim():
    rng = random.Random(SEED)
    sel = pd.read_csv(SEL_CSV, dtype={"date": str})
    mot = pd.read_csv(MOTION_CSV)
    df = sel.merge(mot, on="name", how="inner")
    log(f"{len(df):,} candidates with motion labels")

    is_static = df["motion_components"] == "static"
    statics = df[is_static].sample(frac=1.0, random_state=SEED)
    moving = df[~is_static].sample(frac=1.0, random_state=SEED)
    log(f"  near-static: {len(statics):,} ({len(statics) / len(df):.1%}), moving: {len(moving):,}")

    keep_static = min(len(statics), int(N_FINAL * STATIC_CAP))
    keep_moving = min(len(moving), N_FINAL - keep_static)
    keep_static = min(len(statics), N_FINAL - keep_moving)   # backfill if moving is short
    final = pd.concat([statics.head(keep_static), moving.head(keep_moving)])
    final = final.sample(frac=1.0, random_state=SEED).reset_index(drop=True)
    final.to_csv(FINAL_CSV, index=False)

    log(f"final set: {len(final):,} pairs "
        f"({keep_static} static capped at {STATIC_CAP:.0%}, {keep_moving} moving)")
    comps = collections.Counter()
    for c in final["motion_components"]:
        for part in c.split("+"):
            comps[part] += 1
    log("  motion components: " + ", ".join(
        f"{k} {100 * v / len(final):.1f}%" for k, v in comps.most_common()))
    log(f"  path shape: " + ", ".join(
        f"{k} {100 * v / len(final):.1f}%" for k, v in final['path_shape'].value_counts().items()))
    log(f"  clusters used: {final['cluster'].nunique()}/{N_CLUSTERS}")


# --------------------------------------------------------------------------- encode

def _encode_one(job):
    src, dst = job
    if NATIVE:
        # Remux only — the video stream is copied bit-for-bit, so this is lossless.
        # ffmpeg still parses the source, so a corrupt input fails loudly here.
        cmd = ["ffmpeg", "-v", "error", "-y", "-i", src, "-c:v", "copy", "-an", dst]
    else:
        cmd = ["ffmpeg", "-v", "error", "-y", "-threads", str(ENCODE_THREADS), "-i", src,
               "-vf", f"scale={OUT_W}:-2,crop={OUT_W}:{OUT_H}",
               "-frames:v", str(OUT_FRAMES),
               "-c:v", "libx264", "-crf", "18", "-preset", "medium",
               "-pix_fmt", "yuv420p", "-an", dst]
    p = subprocess.run(cmd, capture_output=True, text=True)
    return dst, p.returncode, p.stderr.strip()[:300]


def stage_encode():
    from concurrent.futures import ThreadPoolExecutor
    os.makedirs(OUT_DIR, exist_ok=True)
    final = pd.read_csv(FINAL_CSV, dtype={"date": str})
    jobs = []
    for r in final.itertuples(index=False):
        jobs.append((f"{SRC_ROOT}/{r.target_src}", f"{OUT_DIR}/{r.name}.mp4"))
        jobs.append((f"{SRC_ROOT}/{r.ref_src}", f"{OUT_DIR}/{r.name}_camera.mp4"))
    jobs = [(s, d) for s, d in jobs if not os.path.exists(d) or os.path.getsize(d) == 0]
    if NATIVE:
        log(f"remuxing {len(jobs):,} clips -> native {NATIVE_CODEC}, {NATIVE_FRAMES} frames, "
            f"source resolution ({ENCODE_JOBS} jobs)")
    else:
        log(f"encoding {len(jobs):,} clips -> {OUT_W}x{OUT_H}, {OUT_FRAMES} frames, h264 "
            f"({ENCODE_JOBS} jobs x {ENCODE_THREADS} threads)")

    fails, t0 = [], time.time()
    with ThreadPoolExecutor(ENCODE_JOBS) as ex:
        for k, (dst, rc, err) in enumerate(ex.map(_encode_one, jobs), 1):
            if rc != 0:
                fails.append((dst, err))
            if k % 500 == 0:
                log(f"  {k}/{len(jobs)} ({time.time() - t0:.0f}s)")
    log(f"encoded in {time.time() - t0:.0f}s, {len(fails)} failures")
    for d, e in fails[:10]:
        log(f"  FAIL {d}: {e}")
    if fails:
        sys.exit(1)


# --------------------------------------------------------------------------- emit

def stage_emit():
    final = pd.read_csv(FINAL_CSV, dtype={"date": str})
    src_dims = {r.name: native_dims(r.date) for r in final.itertuples(index=False)}

    entries, rows = [], []
    for r in final.itertuples(index=False):
        media = f"{r.name}.mp4"
        ref = f"{r.name}_camera.mp4"
        entries.append({"caption": r.caption, "media_path": media, "reference_path": ref})
        w, h = src_dims[r.name]
        rows.append(dict(
            media_path=media, reference_path=ref,
            target_src_path=f"./{r.target_src}", reference_src_path=f"./{r.ref_src}",
            date=r.date, traj_group=r.grp, cam_index=r.cam,
            target_scene_id=r.tsid, ref_scene_id=r.rsid,
            src_width=w, src_height=h,
            motion_components=r.motion_components, n_axes=r.n_axes, path_shape=r.path_shape,
            pan=round(r.pan, 2), tilt=round(r.tilt, 2), roll=round(r.roll, 2),
            zoom=round(r.zoom, 2), straightness=round(r.straightness, 3),
            caption_cluster=r.cluster, caption=r.caption))

    with open(f"{OUT_DIR}/dataset.json", "w") as f:
        json.dump(entries, f, indent=2, ensure_ascii=False)
    pd.DataFrame(rows).to_csv(MAPPING_CSV, index=False)
    log(f"wrote {OUT_DIR}/dataset.json ({len(entries):,} entries)")
    log(f"wrote {MAPPING_CSV} ({len(rows):,} rows)")


# --------------------------------------------------------------------------- verify

def _probe(path):
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
         "stream=width,height,nb_frames,codec_name,pix_fmt", "-of", "csv=p=0", path],
        capture_output=True, text=True).stdout.strip()
    return out


def stage_verify():
    from concurrent.futures import ThreadPoolExecutor
    ok = True

    def check(label, cond, detail=""):
        nonlocal ok
        print(f"  [{'PASS' if cond else 'FAIL'}] {label}{(' — ' + detail) if detail else ''}")
        if not cond:
            ok = False

    ds = json.load(open(f"{OUT_DIR}/dataset.json"))
    mp = pd.read_csv(MAPPING_CSV, dtype={"date": str})
    listing = os.listdir(OUT_DIR)

    print("\n1. counts")
    check(f"{len(listing)} entries in {os.path.basename(OUT_DIR)}",
          len(listing) == 2 * N_FINAL + 1, f"expected {2 * N_FINAL + 1}")
    check(f"{len(ds)} dataset.json records", len(ds) == N_FINAL)
    check("mapping CSV rows match", len(mp) == N_FINAL)
    names = [e["media_path"] for e in ds] + [e["reference_path"] for e in ds]
    check("no duplicate filenames", len(set(names)) == len(names))
    check("all referenced files exist",
          all(os.path.exists(f"{OUT_DIR}/{n}") for n in names))

    print("\n2. schema parity with Depth-Dataset")
    ref_ds = json.load(open("/mnt/dataset/xinyuy/datasets/LTX-2/Depth-Dataset/dataset.json"))
    check("dataset.json keys identical",
          set(ds[0].keys()) == set(ref_ds[0].keys()), str(sorted(ds[0].keys())))
    check("all captions non-empty",
          all(isinstance(e["caption"], str) and e["caption"].strip() for e in ds))

    print(f"\n3. video format (ffprobe, all {len(names)} clips)")
    with ThreadPoolExecutor(48) as ex:
        probes = dict(zip(names, ex.map(lambda n: _probe(f"{OUT_DIR}/{n}"), names)))
    def expected(name):
        # ffprobe csv order is codec_name,width,height,pix_fmt,nb_frames
        if NATIVE:
            w, h = native_dims(name.split("_")[0])
            return f"{NATIVE_CODEC},{w},{h},{NATIVE_PIX_FMT},{NATIVE_FRAMES}"
        return f"h264,{OUT_W},{OUT_H},yuv420p,{OUT_FRAMES}"

    label = (f"all clips are native {NATIVE_CODEC}/{NATIVE_PIX_FMT}/{NATIVE_FRAMES}f at their "
             f"batch resolution" if NATIVE else f"all clips are {expected(names[0])}")
    bad = [n for n, p in probes.items() if p != expected(n)]
    check(label, not bad,
          f"{len(bad)} mismatched, e.g. {bad[:3]} -> {[probes[b] for b in bad[:3]]}" if bad else "")
    if NATIVE:
        dims = collections.Counter(",".join(p.split(",")[1:3]) for p in probes.values())
        print("     resolutions: " + ", ".join(f"{k.replace(',', 'x')} {v}" for k, v in dims.most_common()))
    desync = [e["media_path"] for e in ds
              if probes[e["media_path"]].split(",")[-1] != probes[e["reference_path"]].split(",")[-1]]
    check("target/reference frame counts match (motion sync)", not desync,
          f"{len(desync)} desynced" if desync else "")

    print("\n4. pair correctness")
    same_traj = mp.apply(lambda r: os.path.dirname(r.target_src_path)
                         == os.path.dirname(r.reference_src_path), axis=1)
    check("every pair shares a trajectory dir (same camera motion)", same_traj.all(),
          f"{(~same_traj).sum()} violations")
    diff_scene = (mp["target_scene_id"] != mp["ref_scene_id"])
    check("every pair uses different scenes", diff_scene.all(),
          f"{(~diff_scene).sum()} violations")
    check("source files still exist",
          all(os.path.exists(f"{SRC_ROOT}/{p[2:]}") for p in mp["target_src_path"].head(200)))

    print("\n5. diversity")
    check(f"{mp.groupby(['date', 'traj_group']).ngroups} distinct trajectory groups",
          mp.groupby(["date", "traj_group"]).ngroups == N_FINAL)
    check(f"{mp.groupby(['date', 'target_scene_id']).ngroups} distinct target locations",
          mp.groupby(["date", "target_scene_id"]).ngroups == N_FINAL)

    other = [s for s in SPLITS if s != SPLIT]
    for o in other:
        om = f"{WORK}/camera_dataset_mapping{SPLITS[o]['sfx']}.csv"
        if not os.path.exists(om):
            continue
        om = pd.read_csv(om, dtype={"date": str})
        g_self = set(zip(mp["date"], mp["traj_group"]))
        g_other = set(zip(om["date"], om["traj_group"]))
        loc_self = set(zip(mp["date"], mp["target_scene_id"])) | set(zip(mp["date"], mp["ref_scene_id"]))
        loc_other = set(zip(om["date"], om["target_scene_id"])) | set(zip(om["date"], om["ref_scene_id"]))
        f_self = set(mp["media_path"]) | set(mp["reference_path"])
        f_other = set(om["media_path"]) | set(om["reference_path"])
        check(f"no trajectory group shared with '{o}' split", not (g_self & g_other),
              f"{len(g_self & g_other)} shared")
        check(f"no scene location shared with '{o}' split", not (loc_self & loc_other),
              f"{len(loc_self & loc_other)} shared")
        check(f"no video shared with '{o}' split", not (f_self & f_other),
              f"{len(f_self & f_other)} shared")
    comps = collections.Counter()
    for c in mp["motion_components"]:
        for part in c.split("+"):
            comps[part] += 1
    print("     motion components: " + ", ".join(
        f"{k} {100 * v / len(mp):.1f}%" for k, v in comps.most_common()))
    print("     path shape:        " + ", ".join(
        f"{k} {100 * v / len(mp):.1f}%" for k, v in mp["path_shape"].value_counts().items()))
    print("     date batches:      " + ", ".join(
        f"{k} {v}" for k, v in sorted(mp["date"].value_counts().items())))
    print(f"     caption clusters:  {mp['caption_cluster'].nunique()}/{N_CLUSTERS} used, "
          f"largest {mp['caption_cluster'].value_counts().max()} "
          f"({100 * mp['caption_cluster'].value_counts().max() / len(mp):.1f}%)")

    print("\n6. trajectory-match spot check (30 random pairs re-measured)")
    rng = random.Random(SEED)
    sample = rng.sample(range(len(mp)), 30)
    agree = 0
    checked = 0
    for i in sample:
        r = mp.iloc[i]
        a = measure_trajectory(f"{OUT_DIR}/{r.media_path}")
        b = measure_trajectory(f"{OUT_DIR}/{r.reference_path}")
        if a is None or b is None:
            continue
        checked += 1
        dom = int(np.argmax(np.abs(a[:4])))
        if abs(a[dom]) < 5 or np.sign(a[dom]) == np.sign(b[dom]):
            agree += 1
    check(f"dominant-axis sign agreement {agree}/{checked}", checked and agree / checked >= 0.8)

    total = sum(os.path.getsize(f"{OUT_DIR}/{n}") for n in names)
    print(f"\n  dataset size: {total / 1e9:.1f} GB")
    print(f"\n{'ALL CHECKS PASSED' if ok else 'SOME CHECKS FAILED'}")
    return 0 if ok else 1


# --------------------------------------------------------------------------- main

STAGES = {
    "index": stage_index, "select": stage_select, "measure": stage_measure,
    "trim": stage_trim, "encode": stage_encode, "emit": stage_emit, "verify": stage_verify,
}

if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--stage", default="all",
                    choices=list(STAGES) + ["all"], help="stage to run (default: all)")
    ap.add_argument("--split", default="train", choices=list(SPLITS),
                    help="which dataset to build (default: train). 'val' holds out "
                         "everything 'train' used, so build train first.")
    args = ap.parse_args()
    os.makedirs(WORK, exist_ok=True)
    configure(args.split)
    random.seed(SEED)
    np.random.seed(SEED)
    log(f"split={args.split} -> {OUT_DIR} ({N_FINAL} pairs)")

    order = list(STAGES) if args.stage == "all" else [args.stage]
    if args.split != "train" and "index" in order:
        order.remove("index")   # the index is shared and split-independent
    rc = 0
    for s in order:
        log(f"===== stage: {s} =====")
        rc = STAGES[s]() or 0
    sys.exit(rc)