| |
| """Summarize an immutable dataset manifest and its selected training prefix.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| from collections import Counter, defaultdict |
| from pathlib import Path |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--manifest", required=True, type=Path) |
| parser.add_argument("--training-prefix", type=int, default=2500) |
| parser.add_argument( |
| "--class-component", |
| type=int, |
| default=2, |
| help="zero-based path component containing a class ID", |
| ) |
| parser.add_argument("--output-json", required=True, type=Path) |
| parser.add_argument("--output-markdown", required=True, type=Path) |
| args = parser.parse_args() |
| if args.training_prefix < 1 or args.class_component < 0: |
| raise SystemExit("prefix must be positive and class component non-negative") |
|
|
| manifest_bytes = args.manifest.read_bytes() |
| manifest = json.loads(manifest_bytes) |
| if manifest.get("schema_version") != 1 or not manifest.get("images"): |
| raise SystemExit("unsupported or empty dataset manifest") |
|
|
| split_counts: Counter[str] = Counter() |
| source_counts: dict[str, Counter[str]] = defaultdict(Counter) |
| group_splits: dict[tuple[str, str], str] = {} |
| hashes: set[str] = set() |
| paths: set[str] = set() |
| class_counts: dict[str, Counter[str]] = defaultdict(Counter) |
| train_entries: list[dict[str, object]] = [] |
| for image in manifest["images"]: |
| split = str(image["split"]) |
| source = str(image["source"]) |
| group = str(image["group"]) |
| path = str(image["path"]) |
| digest = str(image["sha256"]) |
| if path in paths or digest in hashes: |
| raise SystemExit(f"duplicate path or content in manifest: {path}") |
| paths.add(path) |
| hashes.add(digest) |
| key = (source, group) |
| previous = group_splits.setdefault(key, split) |
| if previous != split: |
| raise SystemExit(f"group {source}/{group} crosses {previous} and {split}") |
| parts = Path(path).parts |
| if len(parts) <= args.class_component: |
| raise SystemExit(f"{path} has no class component {args.class_component}") |
| class_id = parts[args.class_component] |
| split_counts[split] += 1 |
| source_counts[split][source] += 1 |
| class_counts[split][class_id] += 1 |
| if split == "train": |
| train_entries.append(image) |
|
|
| if len(train_entries) < args.training_prefix: |
| raise SystemExit( |
| f"training split has {len(train_entries)} entries, fewer than prefix " |
| f"{args.training_prefix}" |
| ) |
| prefix_classes = Counter( |
| Path(str(image["path"])).parts[args.class_component] |
| for image in train_entries[: args.training_prefix] |
| ) |
| artifact = { |
| "schema_version": 1, |
| "dataset_name": manifest.get("name"), |
| "manifest_sha256": hashlib.sha256(manifest_bytes).hexdigest(), |
| "provenance": manifest.get("provenance", []), |
| "images": len(manifest["images"]), |
| "unique_content_hashes": len(hashes), |
| "groups": len(group_splits), |
| "split_counts": dict(sorted(split_counts.items())), |
| "source_counts_by_split": { |
| split: dict(sorted(counts.items())) |
| for split, counts in sorted(source_counts.items()) |
| }, |
| "class_counts_by_split": { |
| split: dict(sorted(counts.items())) |
| for split, counts in sorted(class_counts.items()) |
| }, |
| "selected_training_prefix": args.training_prefix, |
| "selected_training_prefix_class_counts": dict(sorted(prefix_classes.items())), |
| } |
| args.output_json.parent.mkdir(parents=True, exist_ok=True) |
| args.output_markdown.parent.mkdir(parents=True, exist_ok=True) |
| args.output_json.write_text(json.dumps(artifact, indent=2) + "\n", encoding="utf-8") |
|
|
| lines = [ |
| "# Dataset summary", |
| "", |
| f"Manifest SHA-256: `{artifact['manifest_sha256']}`.", |
| "", |
| "| Split | Images | Classes | Sources |", |
| "|---|---:|---:|---|", |
| ] |
| for split, count in sorted(split_counts.items()): |
| lines.append( |
| f"| {split} | {count} | {len(class_counts[split])} " |
| f"| {', '.join(f'{source}: {n}' for source, n in sorted(source_counts[split].items()))} |" |
| ) |
| lines.extend( |
| [ |
| "", |
| f"The selected training prefix contains {args.training_prefix} images:", |
| "", |
| "| Class ID | Images |", |
| "|---|---:|", |
| ] |
| ) |
| lines.extend(f"| {name} | {count} |" for name, count in sorted(prefix_classes.items())) |
| args.output_markdown.write_text("\n".join(lines) + "\n", encoding="utf-8") |
| print(f"wrote {args.output_json} and {args.output_markdown}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
