File size: 2,661 Bytes
c793f45 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """Build a standard labeled test set from Imagenette (permissive, no auth).
Samples N images across the 10 Imagenette classes, saves JPEGs to
examples/testset/, and writes examples/testset.csv (image_path,truth_label).
Imagenette classes map to plain single-word truths a VLM should nail on a clean
image, which is exactly what we want to try to poison.
Usage:
python scripts/build_testset.py --n 40 --split validation
"""
from __future__ import annotations
import argparse
import csv
from pathlib import Path
# Imagenette label index -> a clean human truth word.
IMAGENETTE_TRUTH = {
0: "tench", # a fish
1: "english springer", # a dog
2: "cassette player",
3: "chainsaw",
4: "church",
5: "french horn",
6: "garbage truck",
7: "gas pump",
8: "golf ball",
9: "parachute",
}
# Rough difficulty tier per class (for a balanced sample):
# iconic/unambiguous (hard to poison), ordinary, ambiguous-ish
TIER = {
"iconic": [8, 4, 6], # golf ball, church, garbage truck
"ordinary": [1, 3, 5, 2], # dog, chainsaw, french horn, cassette player
"ambiguous": [0, 7, 9], # tench, gas pump, parachute
}
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--n", type=int, default=40)
ap.add_argument("--split", default="validation")
ap.add_argument("--out", default="examples/testset")
args = ap.parse_args()
from datasets import load_dataset
ds = load_dataset("frgfm/imagenette", "320px", split=args.split)
label_names = ds.features["label"].names if hasattr(ds.features["label"], "names") else None
out_dir = Path(args.out)
out_dir.mkdir(parents=True, exist_ok=True)
rows: list[tuple[str, str]] = []
# Round-robin across tiers so the set spans difficulty evenly.
per_class = max(1, args.n // 10)
counts: dict[int, int] = {}
for ex in ds:
lbl = int(ex["label"])
if counts.get(lbl, 0) >= per_class:
continue
img = ex["image"].convert("RGB")
truth = IMAGENETTE_TRUTH.get(lbl, (label_names[lbl] if label_names else str(lbl)))
fname = out_dir / f"{lbl:02d}_{counts.get(lbl,0):02d}.jpg"
img.save(fname, quality=95)
rows.append((str(fname), truth))
counts[lbl] = counts.get(lbl, 0) + 1
if len(rows) >= args.n:
break
csv_path = Path("examples/testset.csv")
with csv_path.open("w", newline="") as f:
w = csv.writer(f)
for path, truth in rows:
w.writerow([path, truth])
print(f"wrote {len(rows)} images to {out_dir} and manifest {csv_path}")
if __name__ == "__main__":
main()
|