scripts/make_splits.py

"""Split the judged pairs into train / unseen-validation.

The paper holds out a 500-pair validation set and reports on 75 of them.  We hold
out the same *proportion* and, critically, make the split **disjoint by video**:
no clip that appears in a validation pair (as reference or target) may appear
anywhere in training.  A naive random split would leak, because the same clip is
reused across many pairs.
"""

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pairs", default="data/pairs.jsonl")
    ap.add_argument("--train_out", default="data/pairs_train.jsonl")
    ap.add_argument("--val_out", default="data/pairs_val.jsonl")
    ap.add_argument("--clips", default="data/annotated.jsonl",
                    help="stage-1 clip set to filter against the held-out videos")
    ap.add_argument("--clips_train_out", default="data/clips_train.jsonl")
    # 0.08 of ~930 judged pairs is 74 validation pairs -- the size the paper
    # reports on.  A larger fraction is actively harmful here: every held-out
    # clip quarantines all the *training* pairs it appears in, so val_frac 0.15
    # costs ~145 train pairs relative to 0.08 for validation we do not need.
    ap.add_argument("--val_frac", type=float, default=0.08)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    rows = [json.loads(l) for l in open(args.pairs)]
    rng = random.Random(args.seed)
    order = list(range(len(rows)))
    rng.shuffle(order)

    # Select validation *pairs* first, then quarantine every clip they touch.
    # (Selecting held-out *videos* first and keeping only pairs fully inside that
    # set is far too restrictive at this dataset size -- it yielded 1 val pair.)
    n_val = max(1, int(len(rows) * args.val_frac))
    val_idx = set(order[:n_val])
    val = [rows[i] for i in sorted(val_idx)]

    val_vids = {r["ref_id"] for r in val} | {r["tgt_id"] for r in val}
    # Train pairs must not touch any held-out clip at all, or physics/appearance
    # from a validation clip leaks into training.
    train = [r for i, r in enumerate(rows)
             if i not in val_idx
             and r["ref_id"] not in val_vids and r["tgt_id"] not in val_vids]
    vids = sorted({r["ref_id"] for r in rows} | {r["tgt_id"] for r in rows})

    for path, data in ((args.train_out, train), (args.val_out, val)):
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            for r in data:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")

    # Stage 1 trains on single clips with self-reference augmentation.  It must
    # not see the held-out videos either, or a validation pair's target has
    # already been reconstructed during training.
    clips = [json.loads(l) for l in open(args.clips)]
    clips_train = [c for c in clips if c["id"] not in val_vids]
    with open(args.clips_train_out, "w") as f:
        for c in clips_train:
            f.write(json.dumps(c, ensure_ascii=False) + "\n")

    dropped = len(rows) - len(train) - len(val)
    print(f"videos: {len(vids)} ({len(val_vids)} held out)")
    print(f"train pairs: {len(train)}")
    print(f"val pairs:   {len(val)}")
    print(f"dropped (cross-split, would leak): {dropped}")
    print(f"stage-1 clips: {len(clips_train)} / {len(clips)}")


if __name__ == "__main__":
    main()