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