from __future__ import annotations import argparse import json from pathlib import Path import random from typing import Any, Iterable from datasets import Dataset, load_dataset from .common import write_json DATASET = "ranjaykrishna/visual_genome" def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument("--output-dir", default="artifacts/vg") p.add_argument("--cache-dir", default="/tmp/yurenh2-worldalign-vg-hf") p.add_argument("--nodes", type=int, default=100_000) p.add_argument("--views", type=int, default=32) p.add_argument("--min-regions", type=int, default=32) p.add_argument("--seed", type=int, default=20260728) p.add_argument( "--with-extra-tiers", action="store_true", help="Also load visible relation/attribute and QA annotations.", ) return p.parse_args() def load_config(name: str, cache_dir: str) -> Dataset: return load_dataset( DATASET, name, split="train", trust_remote_code=True, with_image=False, cache_dir=cache_dir, ) def by_image_id(dataset: Dataset, field: str) -> dict[int, list[dict[str, Any]]]: return { int(image_id): values for image_id, values in zip(dataset["image_id"], dataset[field]) } def first_name(obj: dict[str, Any]) -> str: names = obj.get("names") or [] return str(names[0]).strip() if names else "" def visible_relation_texts( image_id: int, relationships: dict[int, list[dict[str, Any]]], attributes: dict[int, list[dict[str, Any]]], ) -> list[str]: texts: list[str] = [] for rel in relationships.get(image_id, []): subject = first_name(rel["subject"]) predicate = str(rel.get("predicate") or "").strip() object_ = first_name(rel["object"]) if subject and predicate and object_: texts.append(f"{subject} {predicate} {object_}.") for obj in attributes.get(image_id, []): name = first_name(obj) for attribute in obj.get("attributes") or []: attribute = str(attribute).strip() if name and attribute: texts.append(f"{name} is {attribute}.") return texts def qa_texts( image_id: int, questions: dict[int, list[dict[str, Any]]] ) -> list[str]: texts: list[str] = [] for item in questions.get(image_id, []): question = str(item.get("question") or "").strip() answer = str(item.get("answer") or "").strip() if question and answer: texts.append(f"Question: {question} Answer: {answer}") return texts def fixed_mixture( groups: list[list[str]], budget: int, rng: random.Random ) -> list[str]: """Draw a fixed-size, approximately balanced mixture without count leakage.""" groups = [list(dict.fromkeys(group)) for group in groups if group] for group in groups: rng.shuffle(group) selected: list[str] = [] while len(selected) < budget and groups: next_groups: list[list[str]] = [] for group in groups: if group and len(selected) < budget: selected.append(group.pop()) if group: next_groups.append(group) groups = next_groups if not selected: raise ValueError("Cannot construct a text bundle from empty groups") seed_values = selected.copy() while len(selected) < budget: selected.append(rng.choice(seed_values)) rng.shuffle(selected) return selected def write_jsonl(path: Path, records: Iterable[dict[str, Any]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as handle: for record in records: handle.write(json.dumps(record, ensure_ascii=False) + "\n") def main() -> None: args = parse_args() if args.views < 1: raise ValueError("--views must be positive") output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) rng = random.Random(args.seed) region_ds = load_config("region_descriptions_v1.2.0", args.cache_dir) relation_lookup: dict[int, list[dict[str, Any]]] = {} attribute_lookup: dict[int, list[dict[str, Any]]] = {} question_lookup: dict[int, list[dict[str, Any]]] = {} if args.with_extra_tiers: relation_lookup = by_image_id( load_config("relationships_v1.2.0", args.cache_dir), "relationships", ) attribute_lookup = by_image_id( load_config("attributes_v1.2.0", args.cache_dir), "attributes", ) question_lookup = by_image_id( load_config("question_answers_v1.2.0", args.cache_dir), "qas", ) eligible = [ idx for idx, regions in enumerate(region_ds["regions"]) if len(regions) >= args.min_regions ] rng.shuffle(eligible) selected = eligible[: min(args.nodes, len(eligible))] # The two orders and all opaque identifiers are independent. Source IDs # appear only in the vision preprocessing file and private truth file. vision_order = selected.copy() text_order = selected.copy() rng.shuffle(vision_order) rng.shuffle(text_order) vision_opaque = {idx: f"v{rank:07d}" for rank, idx in enumerate(vision_order)} text_opaque = {idx: f"t{rank:07d}" for rank, idx in enumerate(text_order)} vision_records: list[dict[str, Any]] = [] text_records: list[dict[str, Any]] = [] truth: list[dict[str, Any]] = [] for idx in selected: row = region_ds[int(idx)] regions = list(row["regions"]) # Fix the view count and use a deterministic per-node sample. This # removes region-count and ordering side channels. node_rng = random.Random(args.seed ^ int(row["image_id"])) node_rng.shuffle(regions) regions = regions[: args.views] visual_regions = [ { "x": int(region["x"]), "y": int(region["y"]), "width": int(region["width"]), "height": int(region["height"]), } for region in regions ] phrases = [str(region["phrase"]).strip() for region in regions] node_rng.shuffle(phrases) image_id = int(row["image_id"]) v_id = vision_opaque[idx] t_id = text_opaque[idx] vision_records.append( { "node_id": v_id, "source_image_id": image_id, "url": row["url"], "width": int(row["width"]), "height": int(row["height"]), "regions": visual_regions, } ) text_record: dict[str, Any] = { "node_id": t_id, "region_closed": phrases, } if args.with_extra_tiers: relation_text = visible_relation_texts( image_id, relation_lookup, attribute_lookup ) qa_text = qa_texts(image_id, question_lookup) text_record["visible_relations"] = fixed_mixture( [phrases, relation_text], args.views, node_rng ) text_record["qa_expanded"] = fixed_mixture( [phrases, relation_text, qa_text], args.views, node_rng ) text_records.append(text_record) truth.append( { "vision_node_id": v_id, "text_node_id": t_id, "source_image_id": image_id, } ) rng.shuffle(vision_records) rng.shuffle(text_records) rng.shuffle(truth) write_jsonl(output_dir / "vision_nodes.private.jsonl", vision_records) write_jsonl(output_dir / "text_nodes.jsonl", text_records) write_jsonl(output_dir / "ground_truth.private.jsonl", truth) nested_sizes = [ size for size in (5_000, 20_000, 50_000, 100_000) if size <= len(selected) ] if len(selected) not in nested_sizes: nested_sizes.append(len(selected)) manifest = { "dataset": DATASET, "seed": args.seed, "selected_nodes": len(selected), "eligible_nodes": len(eligible), "views_per_node": args.views, "minimum_available_regions": args.min_regions, "extra_closure_tiers": args.with_extra_tiers, "nested_scaling_sizes": sorted(set(nested_sizes)), "training_files": { "vision": "vision_nodes.private.jsonl", "text": "text_nodes.jsonl", }, "evaluation_only": "ground_truth.private.jsonl", "leakage_note": ( "The source image ID and URL are required only by the vision " "feature preprocessor. Alignment training must consume extracted " "features with source metadata removed." ), } write_json(output_dir / "manifest.json", manifest) print( f"Wrote {len(selected):,} hidden-pair nodes to {output_dir}; " f"{len(eligible):,} scenes met the region threshold." ) if __name__ == "__main__": main()