def dump_debug_images(
self,
epoch,
dump_images=True,
num_recons=10,
num_samples=25,
debug_period=10,
unnormalize_images=False,
):
"""
:param epoch:
:param dump_images: Set to False to not dump any images.
:param num_recons:
:param num_samples:
:param debug_period: How often do you dump debug images?
:param unnormalize_images: Should your unnormalize images before
dumping them? Set to True if images are floats in [0, 1].
:return:
"""
if not dump_images or epoch % debug_period != 0:
return
example_obs_batch_np = ptu.get_numpy(self.example_obs_batch)
recon_examples_np = ptu.get_numpy(
self.vae.reconstruct(self.example_obs_batch))
top_row_example = example_obs_batch_np[:num_recons]
bottom_row_recon = np.clip(recon_examples_np, 0, 1)[:num_recons]
recon_vis = combine_images_into_grid(
imgs=list(top_row_example) + list(bottom_row_recon),
imwidth=example_obs_batch_np.shape[2],
imheight=example_obs_batch_np.shape[3],
max_num_cols=len(top_row_example),
image_format='CWH',
unnormalize=unnormalize_images,
)
logdir = logger.get_snapshot_dir()
cv2.imwrite(
osp.join(logdir, '{}_recons.png'.format(epoch)),
cv2.cvtColor(recon_vis, cv2.COLOR_RGB2BGR),
)
raw_samples = ptu.get_numpy(self.vae.sample(num_samples))
vae_samples = np.clip(raw_samples, 0, 1)
vae_sample_vis = combine_images_into_grid(
imgs=list(vae_samples),
imwidth=example_obs_batch_np.shape[2],
imheight=example_obs_batch_np.shape[3],
image_format='CWH',
unnormalize=unnormalize_images,
)
cv2.imwrite(
osp.join(logdir, '{}_vae_samples.png'.format(epoch)),
cv2.cvtColor(vae_sample_vis, cv2.COLOR_RGB2BGR),
)
def visualize_representation(algo, epoch):
if ((epoch < save_period and epoch % initial_save_period == 0)
or epoch % save_period == 0 or epoch >= algo.num_epochs - 1):
logdir = logger.get_snapshot_dir()
filename = osp.join(
logdir,
'obj{obj_id}_sweep_visualization_{epoch}.png'.format(
obj_id=obj_to_sweep, epoch=epoch),
)
visualizations = []
for goal_dict in goal_dicts:
start_img = goal_dict['image_observation']
new_states = goal_dict['new_states']
encoded = encoder.encode(new_states)
# img_format = renderer.output_image_format
images_to_stack = [start_img]
for i in range(encoded.shape[1]):
values = encoded[:, i]
value_image = values.reshape(env_renderer.image_chw[1:])
# TODO: fix hardcoding of CHW
value_img_rgb = np.repeat(value_image[None, :, :],
3,
axis=0)
value_img_rgb = (
(value_img_rgb - value_img_rgb.min()) /
(value_img_rgb.max() - value_img_rgb.min() + 1e-9))
images_to_stack.append(value_img_rgb)
visualizations.append(
combine_images_into_grid(
images_to_stack,
imwidth=renderer.image_chw[2],
imheight=renderer.image_chw[1],
max_num_cols=len(start_states),
pad_length=1,
pad_color=0,
subpad_length=1,
subpad_color=128,
image_format=renderer.output_image_format,
unnormalize=True,
))
final_image = combine_images_into_grid(
visualizations,
imwidth=visualizations[0].shape[1],
imheight=visualizations[0].shape[0],
max_num_cols=3,
image_format='HWC',
pad_length=0,
subpad_length=0,
)
cv2.imwrite(filename, final_image)
print("Saved visualization image to to ", filename)
def save_imgs(imgs, file_path, **kwargs):
imwidth = imgs[0].shape[1]
imheight = imgs[0].shape[2]
imgs = np.clip(imgs, 0, 1)
combined_img = combine_images_into_grid(imgs=list(imgs),
imwidth=imwidth,
imheight=imheight,
**kwargs)
cv2_img = cv2.cvtColor(combined_img, cv2.COLOR_RGB2BGR)
cv2.imwrite(file_path, cv2_img)
def dump_paths(
env,
filename,
paths,
keys,
rows=3,
columns=6,
do_timer=True,
dirname_to_save_images=None,
subdirname="rollouts",
imsize=84, # TODO: automatically set if not provided
imwidth=None,
imheight=None,
num_imgs=3, # how many vertical images we stack per rollout
dump_pickle=False,
grayscale=False,
get_extra_imgs=None,
num_columns_per_rollout=1,
**combine_img_kwargs):
# TODO: merge with `dump_video`
if get_extra_imgs is None:
get_extra_imgs = get_generic_env_imgs
# num_channels = env.vae.input_channels
num_channels = 1 if grayscale else 3
frames = []
imwidth = imwidth or imsize # 500
imheight = imheight or imsize # 300
num_gaps = num_imgs - 1 # 2
H = num_imgs * imheight # imsize
W = imwidth # imsize
if len(paths) < rows * columns:
columns = min(columns, len(paths))
rows = max(min(rows, int(len(paths) / columns)), 1)
else:
rows = min(rows, int(len(paths) / columns))
N = rows * columns
for i in range(N):
start = time.time()
path = paths[i]
l = []
for i_in_path, d in enumerate(path['full_observations']):
imgs = [d[k] for k in keys]
imgs = imgs + get_extra_imgs(path, i_in_path, env)
imgs = imgs[:num_imgs]
l.append(
combine_images_into_grid(imgs,
imwidth,
imheight,
max_num_cols=num_columns_per_rollout,
**combine_img_kwargs))
frames += l
if dirname_to_save_images:
rollout_dir = osp.join(dirname_to_save_images, subdirname, str(i))
os.makedirs(rollout_dir, exist_ok=True)
rollout_frames = frames[-101:]
goal_img = np.flip(rollout_frames[0][:imsize, :imsize, :], 0)
scipy.misc.imsave(rollout_dir + "/goal.png", goal_img)
goal_img = np.flip(rollout_frames[1][:imsize, :imsize, :], 0)
scipy.misc.imsave(rollout_dir + "/z_goal.png", goal_img)
for j in range(0, 101, 1):
img = np.flip(rollout_frames[j][imsize:, :imsize, :], 0)
scipy.misc.imsave(rollout_dir + "/" + str(j) + ".png", img)
if do_timer:
print(i, time.time() - start)
outputdata = reshape_for_video(frames, N, rows, columns, num_channels)
skvideo.io.vwrite(filename, outputdata)
print("Saved video to ", filename)
if dump_pickle:
pickle_filename = filename[:-4] + ".p"
pickle.dump(paths, open(pickle_filename, "wb"))
def dump_video(env,
policy,
filename,
rollout_function,
rows=3,
columns=6,
do_timer=True,
horizon=100,
dirname_to_save_images=None,
subdirname="rollouts",
imsize=84,
get_extra_imgs=None,
grayscale=False,
keys_to_show=None,
num_columns_per_rollout=1,
**combine_img_kwargs):
"""
:param env:
:param policy:
:param filename:
:param rollout_function:
:param rows:
:param columns:
:param pad_length:
:param subpad_color:
:param do_timer:
:param horizon:
:param dirname_to_save_images:
:param subdirname:
:param imsize: TODO: automatically set if not provided
:param get_extra_imgs: A function with type
def get_extra_imgs(
path: List[dict],
index_in_path: int,
env,
) -> List[np.ndarray]:
:param grayscale:
:return:
"""
if get_extra_imgs is None:
get_extra_imgs = get_generic_env_imgs
num_channels = 1 if grayscale else 3
keys_to_show = keys_to_show or ['image_desired_goal', 'image_observation']
frames = []
N = rows * columns
for i in range(N):
start = time.time()
path = rollout_function(
env,
policy,
max_path_length=horizon,
render=False,
)
l = []
for i_in_path, d in enumerate(path['full_observations']):
imgs_to_stack = [d[k] for k in keys_to_show]
imgs_to_stack += get_extra_imgs(path, i_in_path, env)
grid_img = combine_images_into_grid(
imgs_to_stack,
max_num_cols=num_columns_per_rollout,
imwidth=imsize,
imheight=imsize,
unnormalize=True,
**combine_img_kwargs)
l.append(grid_img)
if len(l) < horizon:
frozen_img = l[-1] / 2
l += [frozen_img] * (horizon - len(l))
frames += l
if dirname_to_save_images:
rollout_dir = osp.join(dirname_to_save_images, subdirname, str(i))
os.makedirs(rollout_dir, exist_ok=True)
rollout_frames = frames[-101:]
goal_img = np.flip(rollout_frames[0][:imsize, :imsize, :], 0)
scipy.misc.imsave(rollout_dir + "/goal.png", goal_img)
goal_img = np.flip(rollout_frames[1][:imsize, :imsize, :], 0)
scipy.misc.imsave(rollout_dir + "/z_goal.png", goal_img)
for j in range(0, 101, 1):
img = np.flip(rollout_frames[j][imsize:, :imsize, :], 0)
scipy.misc.imsave(rollout_dir + "/" + str(j) + ".png", img)
if do_timer:
print(i, time.time() - start)
outputdata = reshape_for_video(frames, N, rows, columns, num_channels)
skvideo.io.vwrite(filename, outputdata)
print("Saved video to ", filename)