Spaces:
Running
Running
github-actions[bot]
Deploy from GitHub 1af737237f5d10dfa9cf12849d98f47c3b30ce37 (filtered: code + Dockerfile + README + NOTICES only)
a2f55b2 | """ | |
| Produce stratified train/val/test splits from the master manifest. | |
| Stratification preserves the per-class ratio across splits — important when | |
| the classes are imbalanced (and they always are, eventually). | |
| Determinism: a fixed seed plus sha256 sorting means re-running on the same | |
| manifest produces identical splits. This matters for reproducibility. | |
| Usage | |
| ----- | |
| python scripts/dataset/split.py \ | |
| --manifest data/manifest.csv \ | |
| --out-dir data \ | |
| --val 0.1 --test 0.1 | |
| Generator-aware Stage 3A split: | |
| python scripts/dataset/split.py \ | |
| --manifest data/manifest.csv \ | |
| --out-dir data \ | |
| --stratify-by class-generator | |
| Held-out generator evaluation split: | |
| python scripts/dataset/split.py \ | |
| --manifest data/manifest.csv \ | |
| --out-dir data \ | |
| --stratify-by class-generator \ | |
| --holdout-generator sdxl | |
| Content-disjoint held-out split (recommended for generalisation claims): | |
| python scripts/dataset/split.py \ | |
| --manifest data/manifest.csv \ | |
| --out-dir data/splits_v2 \ | |
| --stratify-by class-generator \ | |
| --holdout-generator sdxl \ | |
| --group-by content-id | |
| Why --group-by content-id exists | |
| -------------------------------- | |
| Output filenames are content-addressed from the (prompt, seed) pair, and every | |
| Stage 3A generator consumed the same `prompts.txt` with the same deterministic | |
| seed derivation. So `<stem>.png` denotes the same prompt and seed in the SDXL, | |
| SD 3.5 and AuraFlow sets. | |
| `--holdout-generator` alone takes *every* row of that generator, which means | |
| the held-out set can depict prompts that a sibling generator already | |
| contributed to training. That is exactly what happened in Stage 3A: all 20,000 | |
| held-out SDXL images shared their (prompt, seed) with an SD 3.5 image in | |
| train/val/test, so the split measured "unseen generator, **seen** content" | |
| while being described as a true generalisation test. | |
| With `--group-by content-id`, a (prompt, seed) identity is assigned to exactly | |
| one partition, and two held-out files are written: | |
| heldout.csv unseen generator AND unseen content (honest test) | |
| heldout_seen_content.csv unseen generator, seen content (legacy-comparable) | |
| Rows of non-held-out generators that fall in the reserved held-out content | |
| partition are dropped from train/val/test and reported, since including them | |
| would reintroduce the leak. | |
| Default remains `--group-by none`, which reproduces the legacy behaviour | |
| byte-for-byte. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import csv | |
| import random | |
| from collections import defaultdict | |
| from pathlib import Path | |
| def _content_id(row: dict[str, str]) -> str: | |
| """Identity of the (prompt, seed) pair behind a content-addressed filename. | |
| The same stem denotes the same prompt and seed across every Stage 3A | |
| generator, so this groups sibling renderings of identical content. | |
| """ | |
| return Path(row["path"]).stem | |
| def _generator_label(row: dict[str, str]) -> str: | |
| """Return the best available generator label for an AI row.""" | |
| return row.get("generator") or row.get("source") or "unknown" | |
| def _split_key(row: dict[str, str], stratify_by: str) -> str: | |
| """Return the grouping key used for stratified splitting.""" | |
| cls = row["class"] | |
| if stratify_by == "class-generator" and cls == "ai_generated": | |
| return f"{cls}:{_generator_label(row)}" | |
| return cls | |
| def _write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, str]]) -> None: | |
| with path.open("w", newline="") as fh: | |
| writer = csv.DictWriter(fh, fieldnames=fieldnames) | |
| writer.writeheader() | |
| writer.writerows(rows) | |
| def _split_content_disjoint( | |
| args: argparse.Namespace, | |
| rows: list[dict[str, str]], | |
| fieldnames: list[str], | |
| holdout_generators: set[str], | |
| ) -> None: | |
| """Partition by (prompt, seed) identity so held-out content is truly unseen.""" | |
| ai_rows = [r for r in rows if r["class"] == "ai_generated"] | |
| real_rows = [r for r in rows if r["class"] != "ai_generated"] | |
| groups: dict[str, list[dict[str, str]]] = defaultdict(list) | |
| for r in ai_rows: | |
| groups[_content_id(r)].append(r) | |
| group_ids = sorted(groups) | |
| rng = random.Random(args.seed) | |
| rng.shuffle(group_ids) | |
| n_groups = len(group_ids) | |
| n_holdout = int(round(n_groups * args.holdout_frac)) if holdout_generators else 0 | |
| holdout_ids = set(group_ids[:n_holdout]) | |
| remaining = group_ids[n_holdout:] | |
| n_test = int(round(len(remaining) * args.test)) | |
| n_val = int(round(len(remaining) * args.val)) | |
| partition = {gid: "test" for gid in remaining[:n_test]} | |
| partition.update({gid: "val" for gid in remaining[n_test : n_test + n_val]}) | |
| partition.update({gid: "train" for gid in remaining[n_test + n_val :]}) | |
| buckets: dict[str, list[dict[str, str]]] = defaultdict(list) | |
| dropped: dict[str, int] = defaultdict(int) | |
| for gid, items in groups.items(): | |
| in_holdout_content = gid in holdout_ids | |
| for r in items: | |
| is_holdout_gen = _generator_label(r) in holdout_generators | |
| if is_holdout_gen: | |
| buckets["heldout" if in_holdout_content else "heldout_seen_content"].append(r) | |
| elif in_holdout_content: | |
| # Including this would put the held-out prompt back into | |
| # training via a sibling generator — the exact Stage 3A leak. | |
| dropped[_generator_label(r)] += 1 | |
| else: | |
| buckets[partition[gid]].append(r) | |
| # Real photos share no content identity with generated images, so they are | |
| # split by the ordinary stratified fractions. | |
| real_sorted = sorted(real_rows, key=lambda r: r["sha256"]) | |
| rng_real = random.Random(args.seed) | |
| rng_real.shuffle(real_sorted) | |
| n = len(real_sorted) | |
| n_test_r = int(round(n * args.test)) | |
| n_val_r = int(round(n * args.val)) | |
| buckets["test"].extend(real_sorted[:n_test_r]) | |
| buckets["val"].extend(real_sorted[n_test_r : n_test_r + n_val_r]) | |
| buckets["train"].extend(real_sorted[n_test_r + n_val_r :]) | |
| print(f"content-id groups: {n_groups}") | |
| print(f" reserved for held-out content: {len(holdout_ids)}") | |
| print(f" train/val/test groups: {len(remaining) - n_test - n_val}/{n_val}/{n_test}") | |
| if dropped: | |
| total_dropped = sum(dropped.values()) | |
| print( | |
| f" dropped {total_dropped} non-held-out rows whose content is " | |
| "reserved (prevents leakage):" | |
| ) | |
| for gen, count in sorted(dropped.items()): | |
| print(f" {gen}: {count}") | |
| args.out_dir.mkdir(parents=True, exist_ok=True) | |
| for name in ("train", "val", "test", "heldout", "heldout_seen_content"): | |
| rs = buckets.get(name, []) | |
| if not rs: | |
| continue | |
| rs.sort(key=lambda r: (r["sha256"], r["path"])) | |
| path = args.out_dir / f"{name}.csv" | |
| _write_csv(path, fieldnames, rs) | |
| counts: dict[str, int] = defaultdict(int) | |
| for r in rs: | |
| key = ( | |
| _generator_label(r) if r["class"] == "ai_generated" else "authentic" | |
| ) | |
| counts[key] += 1 | |
| breakdown = " ".join(f"{k}={v}" for k, v in sorted(counts.items())) | |
| print(f" wrote {path} ({len(rs)} rows) {breakdown}") | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--manifest", type=Path, required=True) | |
| parser.add_argument("--out-dir", type=Path, required=True) | |
| parser.add_argument("--val", type=float, default=0.1) | |
| parser.add_argument("--test", type=float, default=0.1) | |
| parser.add_argument("--seed", type=int, default=0) | |
| parser.add_argument( | |
| "--stratify-by", | |
| choices=["class", "class-generator"], | |
| default="class", | |
| help=( | |
| "Stratification mode. 'class' preserves legacy behavior; " | |
| "'class-generator' additionally balances AI generators." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--holdout-generator", | |
| action="append", | |
| default=[], | |
| help=( | |
| "AI generator to remove from train/val/test and write to " | |
| "heldout.csv. Can be passed multiple times." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--group-by", | |
| choices=["none", "content-id"], | |
| default="none", | |
| help=( | |
| "'none' (default) reproduces legacy behaviour. 'content-id' keeps " | |
| "every generator's rendering of the same (prompt, seed) in one " | |
| "partition, so the held-out split is content-disjoint as well as " | |
| "generator-disjoint." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--holdout-frac", | |
| type=float, | |
| default=0.1, | |
| help=( | |
| "Fraction of content-id groups reserved for the held-out content " | |
| "partition. Only used with --group-by content-id." | |
| ), | |
| ) | |
| args = parser.parse_args() | |
| if args.val + args.test >= 1.0: | |
| raise ValueError("val + test must be < 1.0") | |
| with args.manifest.open() as fh: | |
| reader = csv.DictReader(fh) | |
| rows = list(reader) | |
| fieldnames = reader.fieldnames or [] | |
| holdout_generators = set(args.holdout_generator) | |
| if args.group_by == "content-id": | |
| return _split_content_disjoint(args, rows, fieldnames, holdout_generators) | |
| split_candidates: list[dict[str, str]] = [] | |
| holdout_rows: list[dict[str, str]] = [] | |
| for r in rows: | |
| if r["class"] == "ai_generated" and _generator_label(r) in holdout_generators: | |
| holdout_rows.append(r) | |
| else: | |
| split_candidates.append(r) | |
| if holdout_generators and not holdout_rows: | |
| raise ValueError( | |
| "No rows matched --holdout-generator values: " | |
| f"{sorted(holdout_generators)}" | |
| ) | |
| by_group: dict[str, list[dict[str, str]]] = defaultdict(list) | |
| for r in split_candidates: | |
| by_group[_split_key(r, args.stratify_by)].append(r) | |
| rng = random.Random(args.seed) | |
| train_rows, val_rows, test_rows = [], [], [] | |
| for group, items in sorted(by_group.items()): | |
| # Sort by sha256 for determinism, then shuffle with seeded RNG. | |
| items.sort(key=lambda r: r["sha256"]) | |
| rng.shuffle(items) | |
| n = len(items) | |
| n_test = int(round(n * args.test)) | |
| n_val = int(round(n * args.val)) | |
| test_rows.extend(items[:n_test]) | |
| val_rows.extend(items[n_test : n_test + n_val]) | |
| train_rows.extend(items[n_test + n_val :]) | |
| print( | |
| f" {group}: total={n} " | |
| f"train={n - n_test - n_val} val={n_val} test={n_test}" | |
| ) | |
| args.out_dir.mkdir(parents=True, exist_ok=True) | |
| for name, rs in [("train", train_rows), ("val", val_rows), ("test", test_rows)]: | |
| path = args.out_dir / f"{name}.csv" | |
| _write_csv(path, fieldnames, rs) | |
| print(f" wrote {path} ({len(rs)} rows)") | |
| if holdout_generators: | |
| path = args.out_dir / "heldout.csv" | |
| holdout_rows.sort(key=lambda r: (r["sha256"], r["path"])) | |
| _write_csv(path, fieldnames, holdout_rows) | |
| print( | |
| f" wrote {path} ({len(holdout_rows)} rows) " | |
| f"for held-out generators: {sorted(holdout_generators)}" | |
| ) | |
| if __name__ == "__main__": | |
| main() | |