diff options
Diffstat (limited to 'worldalign/vg_prepare.py')
| -rw-r--r-- | worldalign/vg_prepare.py | 265 |
1 files changed, 265 insertions, 0 deletions
diff --git a/worldalign/vg_prepare.py b/worldalign/vg_prepare.py new file mode 100644 index 0000000..7d553a8 --- /dev/null +++ b/worldalign/vg_prepare.py @@ -0,0 +1,265 @@ +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() |