70 lines
2.7 KiB
Python
70 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
from typing import Any
|
|
|
|
from data.schema import COLORS, OBJECT_IDS, RELATIONS, SHAPES, Fact, SceneObject
|
|
|
|
HELD_OUT_SHAPES = {
|
|
"red": "triangle",
|
|
"blue": "square",
|
|
"green": "star",
|
|
"yellow": "circle",
|
|
}
|
|
|
|
|
|
def _relation(first: SceneObject, second: SceneObject, rng: random.Random) -> Fact:
|
|
relation = rng.choice(RELATIONS)
|
|
if relation == "left_of":
|
|
first = SceneObject(first.object_id, first.color, first.shape, 55, 112)
|
|
second = SceneObject(second.object_id, second.color, second.shape, 169, 112)
|
|
elif relation == "right_of":
|
|
first = SceneObject(first.object_id, first.color, first.shape, 169, 112)
|
|
second = SceneObject(second.object_id, second.color, second.shape, 55, 112)
|
|
elif relation == "above":
|
|
first = SceneObject(first.object_id, first.color, first.shape, 112, 55)
|
|
second = SceneObject(second.object_id, second.color, second.shape, 112, 169)
|
|
else:
|
|
first = SceneObject(first.object_id, first.color, first.shape, 112, 169)
|
|
second = SceneObject(second.object_id, second.color, second.shape, 112, 55)
|
|
return Fact(first.object_id, relation, second.object_id), first, second
|
|
|
|
|
|
def generate_scene(
|
|
scene_id: int, seed: int, num_objects: int = 2, held_out_composition: bool = False
|
|
) -> dict[str, Any]:
|
|
if num_objects != 2:
|
|
raise ValueError("The stage-one MVP currently supports exactly two objects.")
|
|
rng = random.Random((seed * 1_000_003) + scene_id)
|
|
colors = rng.sample(COLORS, k=num_objects)
|
|
if held_out_composition:
|
|
first_shape = HELD_OUT_SHAPES[colors[0]]
|
|
second_choices = [shape for shape in SHAPES if shape != first_shape]
|
|
shapes = [first_shape, rng.choice(second_choices)]
|
|
else:
|
|
shapes = []
|
|
for color in colors:
|
|
choices = [shape for shape in SHAPES if shape != HELD_OUT_SHAPES[color]]
|
|
choices = [shape for shape in choices if shape not in shapes]
|
|
shapes.append(rng.choice(choices))
|
|
objects = [
|
|
SceneObject(OBJECT_IDS[index], colors[index], shapes[index], 0, 0)
|
|
for index in range(num_objects)
|
|
]
|
|
relation, first, second = _relation(objects[0], objects[1], rng)
|
|
objects = [first, second]
|
|
facts = []
|
|
for item in objects:
|
|
facts.extend(
|
|
[
|
|
Fact(item.object_id, "color", item.color),
|
|
Fact(item.object_id, "shape", item.shape),
|
|
]
|
|
)
|
|
facts.append(relation)
|
|
return {
|
|
"scene_id": f"scene-{scene_id:07d}",
|
|
"objects": [item.to_dict() for item in objects],
|
|
"visual_facts": [fact.to_dict() for fact in facts],
|
|
}
|