build_style_transfer_reference_validate.py

#!/usr/bin/env python3
"""Build `style-transfer-reference-validate`: a held-out twin of `style-transfer-reference`.

Same task encoding as the training build — (style reference image + prompt) -> stylized
video — but every clip here is drawn from rows the training set never touched.  Both the
`video` and the `original_video` columns of the training manifest are excluded, so neither
a stylized clip nor its unstylized source can leak across the split.

Selection is `--per-style` clips per style (default 2, i.e. 38 samples across the canonical
19).  Within a style the reference frame is still cut from a *different* clip, so at two
per style the pair simply donates to each other: clip 000's reference is a frame of clip
001 and vice versa.  That keeps the validation samples honest in exactly the way the
training samples are — the reference never shows content from the clip it conditions.

The style taxonomy, prompt-parsing regexes and per-job machinery are imported from
`build_style_transfer_reference` rather than copied, so the two splits cannot drift apart.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor

import pandas as pd

from build_style_transfer_reference import (
    CANONICAL_19,
    DEFAULT_CSV,
    DEFAULT_VIDEO_BASE,
    _NORMALIZE,
    copy_video,
    extract_reference,
    extract_style,
    plan_jobs,
    remap,
    run_pool,
    write_outputs,
)

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_TRAIN = os.path.join(HERE, "style-transfer-reference")
DEFAULT_DST = os.path.join(HERE, "style-transfer-reference-validate")


def load_excluded(train_dir: str) -> set:
    """Every content-hash path the training split consumed, in either role.

    `original_video` is the stylized clip (CSV `video`) and `original_org` its unstylized
    source (CSV `original_video`); a validation row must collide with neither.
    """
    manifest = os.path.join(train_dir, "manifest.csv")
    m = pd.read_csv(manifest)
    used = set(m["original_video"]) | set(m["original_org"])
    print(f"Excluding {len(used):,} video paths used by {os.path.basename(train_dir)} "
          f"({len(m):,} rows)")
    return used


def select_unused_samples(csv_path: str, video_base: str, per_style: int,
                          excluded: set, sample_seed: int) -> pd.DataFrame:
    """Mirror of `select_samples`, with the training split's clips removed up front."""
    print("Loading CSV...")
    data = pd.read_csv(csv_path, sep=";")

    def file_exists(rel_path):
        return os.path.exists(os.path.join(video_base, remap(rel_path)))

    mask = data["video"].apply(file_exists) & data["original_video"].apply(file_exists)
    filtered = data[mask].reset_index(drop=True).copy()
    print(f"Rows with both videos present: {len(filtered):,}")

    filtered["style_raw"] = filtered["prompt"].apply(extract_style)
    filtered["style"] = filtered["style_raw"].apply(lambda r: _NORMALIZE.get(r, r) if r else None)

    gs_common = filtered[
        filtered["style"].isin(CANONICAL_19)
        & filtered["prompt"].str.startswith("Apply")
    ]
    print(f"Forward-only rows across 19 styles: {len(gs_common):,}")

    unseen = gs_common[
        ~gs_common["video"].isin(excluded)
        & ~gs_common["original_video"].isin(excluded)
    ]
    print(f"Held-out rows after removing the training split: {len(unseen):,}")

    # Same shuffled style order + shared-original bookkeeping as the training build, so no
    # unstylized source is reused across two styles within this split either.
    used_originals: set = set()
    selected = []
    order = pd.Series(sorted(CANONICAL_19)).sample(frac=1, random_state=sample_seed).tolist()
    for style in order:
        style_df = unseen[unseen["style"] == style]
        available = style_df[~style_df["original_video"].isin(used_originals)]
        if len(available) < per_style:
            print(f"WARNING: {style!r} only has {len(available)} non-overlapping held-out "
                  f"rows (need {per_style})")
            sampled = available.sample(frac=1, random_state=sample_seed)
        else:
            sampled = available.sample(n=per_style, random_state=sample_seed)
        used_originals.update(sampled["original_video"].tolist())
        selected.append(sampled)

    final_df = pd.concat(selected).reset_index(drop=True)
    assert final_df.groupby("original_video")["style"].nunique().max() == 1, "overlap detected!"
    print(f"\nSelected {len(final_df):,} rows across {final_df['style'].nunique()} styles")
    return final_df


def verify_disjoint(dst: str, excluded: set) -> None:
    written = pd.read_csv(os.path.join(dst, "manifest.csv"))
    leak = set(written["original_video"]) | set(written["original_org"])
    overlap = leak & excluded
    if overlap:
        print(f"\n[FAIL] {len(overlap)} written paths also appear in the training split:")
        for p in sorted(overlap)[:10]:
            print(f"  {p}")
        raise SystemExit(2)
    print(f"\n[ok] all {len(written)} rows disjoint from the training split "
          f"(checked {len(leak)} paths in both roles)")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--csv", default=DEFAULT_CSV)
    ap.add_argument("--video-base", default=DEFAULT_VIDEO_BASE)
    ap.add_argument("--train-dir", default=DEFAULT_TRAIN,
                    help="split whose clips must not reappear here")
    ap.add_argument("--dst", default=DEFAULT_DST)
    ap.add_argument("--per-style", type=int, default=2, help="samples per style")
    ap.add_argument("--sample-seed", type=int, default=7, help="seed for clip selection")
    ap.add_argument("--seed", type=int, default=20260802, help="seed for donor/frame choice")
    ap.add_argument("--ext", default="png", choices=["png", "jpg"])
    ap.add_argument("--jpeg-quality", type=int, default=95)
    ap.add_argument("--workers", type=int, default=min(32, os.cpu_count() or 8))
    ap.add_argument("--no-clean", action="store_true",
                    help="keep existing files in --dst instead of clearing it first")
    args = ap.parse_args()

    os.makedirs(args.dst, exist_ok=True)
    excluded = load_excluded(args.train_dir)
    final_df = select_unused_samples(args.csv, args.video_base, args.per_style,
                                     excluded, args.sample_seed)
    for style, grp in final_df.groupby("style"):
        print(f"  {style:<30} {len(grp):>4}")

    jobs = plan_jobs(final_df, args.video_base, args.ext, args.seed)
    print(f"\nPlanned {len(jobs)} held-out reference-conditioned samples -> {args.dst}")

    if not args.no_clean:
        stale = [f for f in os.listdir(args.dst) if os.path.isfile(os.path.join(args.dst, f))]
        print(f"Clearing {len(stale)} existing files from {args.dst} ...")
        for f in stale:
            os.remove(os.path.join(args.dst, f))

    t0 = time.time()
    print(f"\nCopying stylized clips ({args.workers} threads)...")
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        copied, copy_errors = run_pool(pool, copy_video, jobs, "copy", args.dst)

    jobs = [j for j in jobs if copied.get(j.name, {}).get("status") != "error"]
    ok_names = {j.name for j in jobs}
    jobs = [j for j in jobs if j.donor_name in ok_names and j.donor_video]

    print(f"\nExtracting reference frames ({args.workers} processes)...")
    with ProcessPoolExecutor(max_workers=args.workers) as pool:
        refs, ref_errors = run_pool(pool, extract_reference, jobs, "ref",
                                    args.dst, args.seed, args.jpeg_quality)

    jobs = [j for j in jobs if refs.get(j.name, {}).get("status") != "error"]
    write_outputs(args.dst, jobs, refs)

    errors = copy_errors + ref_errors
    print(f"\nDone in {(time.time()-t0)/60:.1f}m — {len(jobs)} samples written, "
          f"{len(errors)} failed")
    print(f"manifest.csv / dataset.json saved to {args.dst}")
    if errors:
        with open(os.path.join(args.dst, "errors.json"), "w") as fh:
            json.dump(errors, fh, indent=2)
        for e in errors[:10]:
            print(f"  [error] {e['name']}: {e['error']}")

    verify_disjoint(args.dst, excluded)
    return 1 if errors else 0


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