Hi-ToM_Dataset / create_world.py
umwyf's picture
Upload 14 files
acad479
raw
history blame contribute delete
No virus
4.6 kB
import numpy as np
SIZE_TINY = 5
SIZE_SMALL = 10
SIZE_LARGE = 30
SIZE_XLARGE = 50
locations = [
"attic",
"back_yard",
"basement",
"bathroom",
"bedroom",
"cellar",
"closet",
"crawlspace",
"den",
"dining_room",
"front_yard",
"garage",
"garden",
"hall",
"hallway",
"kitchen",
"laundry",
"living_room",
"lounge",
"master_bedroom",
"office",
"pantry",
"patio",
"playroom",
"porch",
"staircase",
"study",
"sunroom",
"TV_room",
"workshop",
]
clothing = [
"belt",
"boots",
"cap",
"coat",
"dress",
"gloves",
"hat",
"jacket",
"jeans",
"pajamas",
"pants",
"raincoat",
"scarf",
"shirt",
"shoes",
"skirt",
"slacks",
"slippers",
"socks",
"stockings",
"suit",
"sweater",
"sweatshirt",
"t-shirt",
"tie",
"trousers",
"underclothes",
"underpants",
"undershirt",
]
fruit = [
"apple",
"banana",
"cherry",
"grapefruit",
"grapes",
"lemon",
"lime",
"melon",
"orange",
"peach",
"pear",
"persimmon",
"pineapple",
"plum",
"strawberry",
"tangerine",
"watermelon",
]
vegetables = [
"asparagus",
"beans",
"broccoli",
"cabbage",
"carrot",
"celery",
"corn",
"cucumber",
"eggplant",
"green_pepper",
"lettuce",
"onion",
"peas",
"potato",
"pumpkin",
"radish",
"spinach",
"sweet_potato",
"tomato",
"turnip",
]
objects = fruit + vegetables
containers = [
"box",
"pantry",
"bathtub",
"envelope",
"drawer",
"bottle",
"cupboard",
"basket",
"crate",
"suitcase",
"bucket",
"container",
"treasure_chest",
]
colors = ['green', 'blue', 'red']
containers = ['_'.join([color, container])
for container in containers
for color in colors]
names = [
"Oliver",
"Ethan",
"Liam",
"Benjamin",
"Lucas",
"Alexander",
"Jacob",
"Mason",
"William",
"Gracie",
"James",
"Logan",
"Owen",
"Noah",
"Carter",
"Nathan",
"Jack",
"Aiden",
"Jackson",
"Jayden",
"Emma",
"Olivia",
"Emily",
"Sophia",
"Ava",
"Chloe",
"Charlotte",
"Abigail",
"Amelia",
"Ella",
"Hannah",
"Isabella",
"Aria",
"Lily",
"Mia",
"Isla",
"Avery",
"Elizabeth",
"Mila",
"Evelyn",
]
assert len(locations) >= SIZE_LARGE
assert len(objects) >= SIZE_LARGE
assert len(containers) >= SIZE_LARGE
assert len(names) >= SIZE_LARGE
def write_world(filepath, locs, objs, conts, nams):
with open(filepath, 'w') as f:
f.write('# locations\n')
for loc in locs:
f.write('\n')
f.write('create %s\n' % loc)
f.write('set %s is_thing\n' % loc)
f.write('set %s is_location\n' % loc)
f.write('\n')
f.write('# objects\n')
for obj in objs:
f.write('\n')
f.write('create %s\n' % obj)
f.write('set %s is_thing\n' % obj)
f.write('set %s is_gettable\n' % obj)
f.write('\n')
f.write('# containers\n')
for cont in conts:
f.write('\n')
f.write('create %s\n' % cont)
f.write('set %s is_thing\n' % cont)
f.write('set %s is_container\n' % cont)
f.write('\n')
f.write('# actors\n')
for nam in nams:
f.write('\n')
f.write('create %s\n' % nam)
f.write('set %s is_actor\n' % nam)
f.write('set %s is_god\n' % nam)
write_world('world_tiny.txt',
np.random.choice(locations, SIZE_TINY, replace=False),
np.random.choice(objects, SIZE_TINY, replace=False),
np.random.choice(containers, SIZE_TINY, replace=False),
np.random.choice(names, SIZE_TINY, replace=False))
write_world('world_small.txt',
np.random.choice(locations, SIZE_SMALL, replace=False),
np.random.choice(objects, SIZE_SMALL, replace=False),
np.random.choice(containers, SIZE_SMALL, replace=False),
np.random.choice(names, SIZE_SMALL, replace=False))
write_world('world_large.txt',
np.random.choice(locations, SIZE_LARGE, replace=False),
np.random.choice(objects, SIZE_LARGE, replace=False),
np.random.choice(containers, SIZE_LARGE, replace=False),
np.random.choice(names, SIZE_LARGE, replace=False))