def observations_to_image(observation: Dict, info: Dict) -> np.ndarray:
r"""Generate image of single frame from observation and info
returned from a single environment step().
Args:
observation: observation returned from an environment step().
info: info returned from an environment step().
Returns:
generated image of a single frame.
"""
egocentric_view_l: List[np.ndarray] = []
if 'high_rgb' in observation:
rgb = observation["high_rgb"]
if not isinstance(rgb, np.ndarray):
rgb = rgb.cpu().numpy()
egocentric_view_l.append(rgb)
else:
if "rgb" in observation:
rgb = observation["rgb"]
if not isinstance(rgb, np.ndarray):
rgb = rgb.cpu().numpy()
egocentric_view_l.append(rgb)
# draw depth map if observation has depth info
if "depth" in observation:
depth_map = observation["depth"].squeeze() * 255.0
if not isinstance(depth_map, np.ndarray):
depth_map = depth_map.cpu().numpy()
depth_map = depth_map.astype(np.uint8)
depth_map = np.stack([depth_map for _ in range(3)], axis=2)
egocentric_view_l.append(depth_map)
# add image goal if observation has image_goal info
if "imagegoal" in observation:
rgb = observation["imagegoal"]
if not isinstance(rgb, np.ndarray):
rgb = rgb.cpu().numpy()
egocentric_view_l.append(rgb)
assert (len(egocentric_view_l) >
0), "Expected at least one visual sensor enabled."
egocentric_view = np.concatenate(egocentric_view_l, axis=1)
# draw collision
if "collisions" in info and info["collisions"]["is_collision"]:
egocentric_view = draw_collision(egocentric_view)
frame = egocentric_view
if "top_down_map" in info:
top_down_map = maps.colorize_draw_agent_and_fit_to_height(
info["top_down_map"], egocentric_view.shape[0])
frame = np.concatenate((egocentric_view, top_down_map), axis=1)
return frame
def observations_to_image(observation: Dict, info: Dict) -> np.ndarray:
r"""Generate image of single frame from observation and info
returned from a single environment step().
Args:
observation: observation returned from an environment step().
info: info returned from an environment step().
Returns:
generated image of a single frame.
"""
render_obs_images: List[np.ndarray] = []
for sensor_name in observation:
if "rgb" in sensor_name:
rgb = observation[sensor_name]
if not isinstance(rgb, np.ndarray):
rgb = rgb.cpu().numpy()
render_obs_images.append(rgb)
elif "depth" in sensor_name:
depth_map = observation[sensor_name].squeeze() * 255.0
if not isinstance(depth_map, np.ndarray):
depth_map = depth_map.cpu().numpy()
depth_map = depth_map.astype(np.uint8)
depth_map = np.stack([depth_map for _ in range(3)], axis=2)
render_obs_images.append(depth_map)
# add image goal if observation has image_goal info
if "imagegoal" in observation:
rgb = observation["imagegoal"]
if not isinstance(rgb, np.ndarray):
rgb = rgb.cpu().numpy()
render_obs_images.append(rgb)
assert (
len(render_obs_images) > 0
), "Expected at least one visual sensor enabled."
shapes_are_equal = len(set(x.shape for x in render_obs_images)) == 1
if not shapes_are_equal:
render_frame = tile_images(render_obs_images)
else:
render_frame = np.concatenate(render_obs_images, axis=1)
# draw collision
if "collisions" in info and info["collisions"]["is_collision"]:
render_frame = draw_collision(render_frame)
if "top_down_map" in info:
top_down_map = maps.colorize_draw_agent_and_fit_to_height(
info["top_down_map"], render_frame.shape[0]
)
render_frame = np.concatenate((render_frame, top_down_map), axis=1)
return render_frame
def draw_top_down_map(info, output_size):
return maps.colorize_draw_agent_and_fit_to_height(info["top_down_map"],
output_size)