def _get_stored_demos(self, amount: int, image_paths: bool):
task_root = join(self._dataset_root, self._task.get_name())
if not exists(task_root):
raise RuntimeError("Can't find the demos for %s at: %s" %
(self._task.get_name(), task_root))
# Sample an amount of examples for the variation of this task
examples_path = join(task_root,
VARIATIONS_FOLDER % self._variation_number,
EPISODES_FOLDER)
examples = listdir(examples_path)
if amount == -1:
amount = len(examples)
if amount > len(examples):
raise RuntimeError(
'You asked for %d examples, but only %d were available.' %
(amount, len(examples)))
selected_examples = np.random.choice(examples, amount, replace=False)
obs_config = self._obs_config
# Process these examples (e.g. loading observations)
demos = []
for example in selected_examples:
example_path = join(examples_path, example)
with open(join(example_path, LOW_DIM_PICKLE), 'rb') as f:
obs = pickle.load(f)
l_sh_rgb_f = join(example_path, LEFT_SHOULDER_RGB_FOLDER)
l_sh_depth_f = join(example_path, LEFT_SHOULDER_DEPTH_FOLDER)
l_sh_mask_f = join(example_path, LEFT_SHOULDER_MASK_FOLDER)
r_sh_rgb_f = join(example_path, RIGHT_SHOULDER_RGB_FOLDER)
r_sh_depth_f = join(example_path, RIGHT_SHOULDER_DEPTH_FOLDER)
r_sh_mask_f = join(example_path, RIGHT_SHOULDER_MASK_FOLDER)
wrist_rgb_f = join(example_path, WRIST_RGB_FOLDER)
wrist_depth_f = join(example_path, WRIST_DEPTH_FOLDER)
wrist_mask_f = join(example_path, WRIST_MASK_FOLDER)
num_steps = len(obs)
if not (num_steps == len(listdir(l_sh_rgb_f)) == len(
listdir(l_sh_depth_f)) == len(listdir(r_sh_rgb_f)) == len(
listdir(r_sh_depth_f)) == len(listdir(wrist_rgb_f)) ==
len(listdir(wrist_depth_f))):
raise RuntimeError('Broken dataset assumption')
for i in range(num_steps):
si = IMAGE_FORMAT % i
if obs_config.left_shoulder_camera.rgb:
obs[i].left_shoulder_rgb = join(l_sh_rgb_f, si)
if obs_config.left_shoulder_camera.depth:
obs[i].left_shoulder_depth = join(l_sh_depth_f, si)
if obs_config.left_shoulder_camera.mask:
obs[i].left_shoulder_mask = join(l_sh_mask_f, si)
if obs_config.right_shoulder_camera.rgb:
obs[i].right_shoulder_rgb = join(r_sh_rgb_f, si)
if obs_config.right_shoulder_camera.depth:
obs[i].right_shoulder_depth = join(r_sh_depth_f, si)
if obs_config.right_shoulder_camera.depth:
obs[i].right_shoulder_mask = join(r_sh_mask_f, si)
if obs_config.wrist_camera.rgb:
obs[i].wrist_rgb = join(wrist_rgb_f, si)
if obs_config.wrist_camera.depth:
obs[i].wrist_depth = join(wrist_depth_f, si)
if obs_config.wrist_camera.depth:
obs[i].wrist_mask = join(wrist_mask_f, si)
if not image_paths:
for i in range(num_steps):
if obs_config.left_shoulder_camera.rgb:
obs[i].left_shoulder_rgb = np.array(
self._resize_if_needed(
Image.open(obs[i].left_shoulder_rgb),
obs_config.left_shoulder_camera.image_size))
if obs_config.left_shoulder_camera.depth:
obs[i].left_shoulder_depth = image_to_float_array(
self._resize_if_needed(
Image.open(obs[i].left_shoulder_depth),
obs_config.left_shoulder_camera.image_size),
DEPTH_SCALE)
if obs_config.left_shoulder_camera.mask:
obs[i].left_shoulder_mask = np.array(
self._resize_if_needed(
Image.open(obs[i].left_shoulder_mask),
obs_config.left_shoulder_camera.image_size))
if obs_config.right_shoulder_camera.rgb:
obs[i].right_shoulder_rgb = np.array(
self._resize_if_needed(
Image.open(obs[i].right_shoulder_rgb),
obs_config.right_shoulder_camera.image_size))
if obs_config.right_shoulder_camera.depth:
obs[i].right_shoulder_depth = image_to_float_array(
self._resize_if_needed(
Image.open(obs[i].right_shoulder_depth),
obs_config.right_shoulder_camera.image_size),
DEPTH_SCALE)
if obs_config.right_shoulder_camera.mask:
obs[i].right_shoulder_mask = np.array(
self._resize_if_needed(
Image.open(obs[i].right_shoulder_mask),
obs_config.right_shoulder_camera.image_size))
if obs_config.wrist_camera.rgb:
obs[i].wrist_rgb = np.array(
self._resize_if_needed(
Image.open(obs[i].wrist_rgb),
obs_config.wrist_camera.image_size))
if obs_config.wrist_camera.depth:
obs[i].wrist_depth = image_to_float_array(
self._resize_if_needed(
Image.open(obs[i].wrist_depth),
obs_config.wrist_camera.image_size),
DEPTH_SCALE)
if obs_config.wrist_camera.mask:
obs[i].wrist_mask = np.array(
self._resize_if_needed(
Image.open(obs[i].wrist_mask),
obs_config.wrist_camera.image_size))
demos.append(obs)
return demos
def get_stored_demos(amount: int, image_paths: bool, dataset_root: str,
variation_number: int, task_name: str,
obs_config: ObservationConfig) -> List[Demo]:
task_root = join(dataset_root, task_name)
if not exists(task_root):
raise RuntimeError("Can't find the demos for %s at: %s" %
(task_name, task_root))
# Sample an amount of examples for the variation of this task
examples_path = join(task_root, VARIATIONS_FOLDER % variation_number,
EPISODES_FOLDER)
examples = listdir(examples_path)
if amount == -1:
amount = len(examples)
if amount > len(examples):
raise RuntimeError(
'You asked for %d examples, but only %d were available.' %
(amount, len(examples)))
selected_examples = np.random.choice(examples, amount, replace=False)
# Process these examples (e.g. loading observations)
demos = []
for example in selected_examples:
example_path = join(examples_path, example)
with open(join(example_path, LOW_DIM_PICKLE), 'rb') as f:
obs = pickle.load(f)
l_sh_rgb_f = join(example_path, LEFT_SHOULDER_RGB_FOLDER)
l_sh_depth_f = join(example_path, LEFT_SHOULDER_DEPTH_FOLDER)
l_sh_mask_f = join(example_path, LEFT_SHOULDER_MASK_FOLDER)
r_sh_rgb_f = join(example_path, RIGHT_SHOULDER_RGB_FOLDER)
r_sh_depth_f = join(example_path, RIGHT_SHOULDER_DEPTH_FOLDER)
r_sh_mask_f = join(example_path, RIGHT_SHOULDER_MASK_FOLDER)
wrist_rgb_f = join(example_path, WRIST_RGB_FOLDER)
wrist_depth_f = join(example_path, WRIST_DEPTH_FOLDER)
wrist_mask_f = join(example_path, WRIST_MASK_FOLDER)
front_rgb_f = join(example_path, FRONT_RGB_FOLDER)
front_depth_f = join(example_path, FRONT_DEPTH_FOLDER)
front_mask_f = join(example_path, FRONT_MASK_FOLDER)
num_steps = len(obs)
if not (num_steps == len(listdir(l_sh_rgb_f)) == len(
listdir(l_sh_depth_f)) == len(listdir(r_sh_rgb_f)) == len(
listdir(r_sh_depth_f)) == len(listdir(wrist_rgb_f)) == len(
listdir(wrist_depth_f)) == len(listdir(front_rgb_f)) ==
len(listdir(front_depth_f))):
raise RuntimeError('Broken dataset assumption')
for i in range(num_steps):
si = IMAGE_FORMAT % i
if obs_config.left_shoulder_camera.rgb:
obs[i].left_shoulder_rgb = join(l_sh_rgb_f, si)
if obs_config.left_shoulder_camera.depth:
obs[i].left_shoulder_depth = join(l_sh_depth_f, si)
if obs_config.left_shoulder_camera.mask:
obs[i].left_shoulder_mask = join(l_sh_mask_f, si)
if obs_config.right_shoulder_camera.rgb:
obs[i].right_shoulder_rgb = join(r_sh_rgb_f, si)
if obs_config.right_shoulder_camera.depth:
obs[i].right_shoulder_depth = join(r_sh_depth_f, si)
if obs_config.right_shoulder_camera.mask:
obs[i].right_shoulder_mask = join(r_sh_mask_f, si)
if obs_config.wrist_camera.rgb:
obs[i].wrist_rgb = join(wrist_rgb_f, si)
if obs_config.wrist_camera.depth:
obs[i].wrist_depth = join(wrist_depth_f, si)
if obs_config.wrist_camera.mask:
obs[i].wrist_mask = join(wrist_mask_f, si)
if obs_config.front_camera.rgb:
obs[i].front_rgb = join(front_rgb_f, si)
if obs_config.front_camera.depth:
obs[i].front_depth = join(front_depth_f, si)
if obs_config.front_camera.mask:
obs[i].front_mask = join(front_mask_f, si)
# Remove low dim info if necessary
if not obs_config.joint_velocities:
obs[i].joint_velocities = None
if not obs_config.joint_positions:
obs[i].joint_positions = None
if not obs_config.joint_forces:
obs[i].joint_forces = None
if not obs_config.gripper_open:
obs[i].gripper_open = None
if not obs_config.gripper_pose:
obs[i].gripper_pose = None
if not obs_config.gripper_joint_positions:
obs[i].gripper_joint_positions = None
if not obs_config.gripper_touch_forces:
obs[i].gripper_touch_forces = None
if not obs_config.task_low_dim_state:
obs[i].task_low_dim_state = None
if not image_paths:
for i in range(num_steps):
if obs_config.left_shoulder_camera.rgb:
obs[i].left_shoulder_rgb = np.array(
_resize_if_needed(
Image.open(obs[i].left_shoulder_rgb),
obs_config.left_shoulder_camera.image_size))
if obs_config.right_shoulder_camera.rgb:
obs[i].right_shoulder_rgb = np.array(
_resize_if_needed(
Image.open(obs[i].right_shoulder_rgb),
obs_config.right_shoulder_camera.image_size))
if obs_config.wrist_camera.rgb:
obs[i].wrist_rgb = np.array(
_resize_if_needed(Image.open(obs[i].wrist_rgb),
obs_config.wrist_camera.image_size))
if obs_config.front_camera.rgb:
obs[i].front_rgb = np.array(
_resize_if_needed(Image.open(obs[i].front_rgb),
obs_config.front_camera.image_size))
if obs_config.left_shoulder_camera.depth:
obs[i].left_shoulder_depth = image_to_float_array(
_resize_if_needed(
Image.open(obs[i].left_shoulder_depth),
obs_config.left_shoulder_camera.image_size),
DEPTH_SCALE)
if obs_config.right_shoulder_camera.depth:
obs[i].right_shoulder_depth = image_to_float_array(
_resize_if_needed(
Image.open(obs[i].right_shoulder_depth),
obs_config.right_shoulder_camera.image_size),
DEPTH_SCALE)
if obs_config.wrist_camera.depth:
obs[i].wrist_depth = image_to_float_array(
_resize_if_needed(Image.open(obs[i].wrist_depth),
obs_config.wrist_camera.image_size),
DEPTH_SCALE)
if obs_config.front_camera.depth:
obs[i].front_depth = image_to_float_array(
_resize_if_needed(Image.open(obs[i].front_depth),
obs_config.front_camera.image_size),
DEPTH_SCALE)
# Masks are stored as coded RGB images.
# Here we transform them into 1 channel handles.
if obs_config.left_shoulder_camera.mask:
obs[i].left_shoulder_mask = rgb_handles_to_mask(
np.array(
_resize_if_needed(
Image.open(obs[i].left_shoulder_mask),
obs_config.left_shoulder_camera.image_size)))
if obs_config.right_shoulder_camera.mask:
obs[i].right_shoulder_mask = rgb_handles_to_mask(
np.array(
_resize_if_needed(
Image.open(obs[i].right_shoulder_mask),
obs_config.right_shoulder_camera.image_size)))
if obs_config.wrist_camera.mask:
obs[i].wrist_mask = rgb_handles_to_mask(
np.array(
_resize_if_needed(
Image.open(obs[i].wrist_mask),
obs_config.wrist_camera.image_size)))
if obs_config.front_camera.mask:
obs[i].front_mask = rgb_handles_to_mask(
np.array(
_resize_if_needed(
Image.open(obs[i].front_mask),
obs_config.front_camera.image_size)))
demos.append(obs)
return demos