| |
| """Prepare deterministic train/validation relation pools from official CLEVR.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import random |
| import zipfile |
| from collections import Counter |
| from datetime import datetime, timezone |
| from pathlib import Path |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| "--archive", |
| type=Path, |
| default=(Path(__file__).resolve().parents[2] / "datasets/raw/clevr/CLEVR_v1.0.zip"), |
| ) |
| parser.add_argument( |
| "--output-dir", |
| type=Path, |
| default=Path("./datasets/clevr_relations"), |
| ) |
| parser.add_argument("--train-pool", type=int, default=1024) |
| parser.add_argument("--validation-pool", type=int, default=256) |
| parser.add_argument("--seed", type=int, default=20260829) |
| return parser.parse_args() |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for chunk in iter(lambda: handle.read(16 << 20), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def read_questions(archive: zipfile.ZipFile, split: str) -> list[dict]: |
| member = f"CLEVR_v1.0/questions/CLEVR_{split}_questions.json" |
| with archive.open(member) as handle: |
| payload = json.load(handle) |
| return payload["questions"] |
|
|
|
|
| def has_relation_program(sample: dict) -> bool: |
| return any(step.get("function") == "relate" for step in sample.get("program", [])) |
|
|
|
|
| def choose_pool(rows: list[dict], count: int, seed: int) -> list[dict]: |
| eligible = [row for row in rows if has_relation_program(row)] |
| order = list(range(len(eligible))) |
| random.Random(seed).shuffle(order) |
| return [eligible[index] for index in order[: min(count, len(order))]] |
|
|
|
|
| def materialize( |
| archive: zipfile.ZipFile, |
| rows: list[dict], |
| split: str, |
| output_dir: Path, |
| ) -> list[dict]: |
| image_dir = output_dir / "images" / split |
| image_dir.mkdir(parents=True, exist_ok=True) |
| prepared = [] |
| for row in rows: |
| filename = row["image_filename"] |
| target = image_dir / filename |
| if not target.exists(): |
| member = f"CLEVR_v1.0/images/{split}/{filename}" |
| with archive.open(member) as source, target.open("wb") as sink: |
| while chunk := source.read(1 << 20): |
| sink.write(chunk) |
| functions = [step.get("function") for step in row.get("program", [])] |
| prepared.append( |
| { |
| "split": split, |
| "question_index": int(row["question_index"]), |
| "image_index": int(row["image_index"]), |
| "question_family_index": int(row["question_family_index"]), |
| "question": row["question"], |
| "answer": str(row["answer"]), |
| "image_filename": filename, |
| "image_path": str(target), |
| "program_functions": functions, |
| "relation_count": functions.count("relate"), |
| } |
| ) |
| return prepared |
|
|
|
|
| def summarize(rows: list[dict]) -> dict: |
| return { |
| "rows": len(rows), |
| "unique_images": len({row["image_index"] for row in rows}), |
| "answer_distribution": dict(Counter(row["answer"] for row in rows)), |
| "relation_count_distribution": dict( |
| Counter(str(row["relation_count"]) for row in rows) |
| ), |
| "question_family_count": len( |
| {row["question_family_index"] for row in rows} |
| ), |
| } |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| if not args.archive.exists(): |
| raise FileNotFoundError(args.archive) |
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| with zipfile.ZipFile(args.archive) as archive: |
| train_all = read_questions(archive, "train") |
| validation_all = read_questions(archive, "val") |
| train_selected = choose_pool(train_all, args.train_pool, args.seed) |
| validation_selected = choose_pool( |
| validation_all, args.validation_pool, args.seed + 1 |
| ) |
| train_rows = materialize(archive, train_selected, "train", args.output_dir) |
| validation_rows = materialize( |
| archive, validation_selected, "val", args.output_dir |
| ) |
|
|
| train_path = args.output_dir / "train_relation_pool.json" |
| validation_path = args.output_dir / "validation_relation_pool.json" |
| train_path.write_text( |
| json.dumps(train_rows, ensure_ascii=False, indent=2), encoding="utf-8" |
| ) |
| validation_path.write_text( |
| json.dumps(validation_rows, ensure_ascii=False, indent=2), encoding="utf-8" |
| ) |
| manifest = { |
| "dataset": "CLEVR v1.0", |
| "source": "https://dl.fbaipublicfiles.com/clevr/CLEVR_v1.0.zip", |
| "license": "CC BY 4.0", |
| "archive": str(args.archive), |
| "archive_bytes": args.archive.stat().st_size, |
| "archive_sha256": sha256_file(args.archive), |
| "selection": "program contains at least one relate operation", |
| "seed": args.seed, |
| "train": summarize(train_rows), |
| "validation": summarize(validation_rows), |
| "test_rows_read": 0, |
| "prepared_at": datetime.now(timezone.utc).isoformat(), |
| } |
| manifest_path = args.output_dir / "dataset_manifest.json" |
| manifest_path.write_text( |
| json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8" |
| ) |
| print(json.dumps(manifest, ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
