80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
from data.generate_rules import add_rule
|
|
from data.generate_scenes import generate_scene
|
|
from data.render_scene import render_scene
|
|
|
|
|
|
def _write_split(
|
|
name: str,
|
|
size: int,
|
|
start: int,
|
|
root: Path,
|
|
seed: int,
|
|
image_size: int,
|
|
held_out: bool,
|
|
) -> None:
|
|
path = root / f"{name}.jsonl"
|
|
with path.open("w", encoding="utf-8") as handle:
|
|
for scene_number in range(start, start + size):
|
|
scene = generate_scene(
|
|
scene_number, seed, held_out_composition=held_out
|
|
)
|
|
record = add_rule(scene, seed, held_out=held_out)
|
|
relative_image = Path("images") / f"{record['scene_id']}.png"
|
|
render_scene(record, root / relative_image, image_size=image_size)
|
|
record["image"] = str(relative_image)
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
|
|
|
|
def build(config_path: Path, small: bool = False, reuse_existing: bool = False) -> Path:
|
|
config = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
|
root = Path(config["data"]["root"])
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
sizes = {
|
|
"train": config["data"]["train_size"],
|
|
"validation": config["data"]["validation_size"],
|
|
"test_id": config["data"]["test_id_size"],
|
|
"test_ood": config["data"]["test_ood_size"],
|
|
}
|
|
if small:
|
|
sizes = {"train": 64, "validation": 16, "test_id": 16, "test_ood": 16}
|
|
if reuse_existing and all(
|
|
(root / f"{split}.jsonl").exists()
|
|
and sum(1 for _ in (root / f"{split}.jsonl").open(encoding="utf-8")) == size
|
|
for split, size in sizes.items()
|
|
):
|
|
return root
|
|
cursor = 0
|
|
for split, size in sizes.items():
|
|
_write_split(
|
|
split,
|
|
size,
|
|
cursor,
|
|
root,
|
|
config["seed"],
|
|
config["data"]["image_size"],
|
|
held_out=split == "test_ood",
|
|
)
|
|
cursor += size
|
|
return root
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml"))
|
|
parser.add_argument("--small", action="store_true")
|
|
parser.add_argument("--reuse-existing", action="store_true")
|
|
args = parser.parse_args()
|
|
print(build(args.config, small=args.small, reuse_existing=args.reuse_existing))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|