def take_pics():
P.initialize_experiment()
train_i, dev_i, test_i, _ = get_all_instructions()
all_instructions = {**train_i, **dev_i, **test_i}
save_dir = paths.get_env_image_path(0)
os.makedirs(os.path.dirname(save_dir), exist_ok=True)
keylist = list(all_instructions.keys())
envs = [PomdpInterface(instance_id=i) for i in range(0, NUM_WORKERS)]
env_id_splits = [[] for _ in range(NUM_WORKERS)]
keylist = [6825]
for i, key in enumerate(keylist):
env_id_splits[i % NUM_WORKERS].append(key)
time.sleep(1.0)
for i in range(len(keylist)):
d = False
# For each worker, start the correct env
for w in range(NUM_WORKERS):
if i >= len(env_id_splits[w]):
continue
env_id = env_id_splits[w][i]
# FIXME: :This assumes that there is only 1 instruction set per env!
fname = paths.get_env_image_path(env_id)
if os.path.isfile(fname):
print("Img exists: " + fname)
continue
d = True
instruction_set = all_instructions[env_id][0]
envs[w].set_environment(env_id, instruction_set["instructions"], fast=True)
print("setting env on worker " + str(w) + " iter " + str(i) + " env_id: " + str(env_id))
# Then for each worker, take a picture and save it
if d:
time.sleep(0.1)
for w in range(NUM_WORKERS):
if i >= len(env_id_splits[w]):
continue
env_id = env_id_splits[w][i]
fname = paths.get_env_image_path(env_id)
if os.path.isfile(fname):
print("Img exists: " + fname)
continue
envs[w].snap_birdseye(fast=True, small_env=SMALL_ENV)
image = envs[w].snap_birdseye(fast=True, small_env=SMALL_ENV)
image = np.flip(image, 0)
imsave(fname, image)
print("saving pic on worker " + str(w) + " iter " + str(i) + " env_id: " + str(env_id))
def load_env_img(env_id, width=None, height=None):
img = imread(paths.get_env_image_path(env_id))
if width is not None:
img = transform.resize(
img, [width, height], mode="constant")
return np.array(img)
def save_env_image(self, env_id, filename, folder):
if folder != "":
os.makedirs(folder, exist_ok=True)
try:
shutil.copy(get_env_image_path(env_id),
os.path.join(folder, filename))
except Exception as e:
print("Error saving env image!")
print(e)
def load_env_img(env_id,
width=None,
height=None,
real_drone=False,
origin_bottom_left=False,
flipdiag=False,
alpha=False,
directory_override=None):
img = imread(paths.get_env_image_path(env_id, real_drone=real_drone, dir_override=directory_override))
if width is not None:
img = transform.resize(
img, [width, height], mode="constant")
# Flip vertically to put origin on bottom-left corner
if origin_bottom_left:
img = np.flip(img, 0)
if flipdiag:
img = img.transpose([1, 0, 2])
if alpha:
img = np.concatenate([img, np.ones(img[:,:,0:1].shape)], axis=2)
return np.array(img)
def save_env_img(env_id, img, real_drone=False):
path = paths.get_env_image_path(env_id, real_drone=real_drone)
dirname = os.path.dirname(path)
os.makedirs(dirname, exist_ok=True)
imsave(path, img)
print(f"Saved env image for env {env_id} in {path}")