scripts/generate_maskrcnn.py [435:562]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - return panorama_frames, camera_info class FailedReplay(Exception): def __init__(self, message): super(FailedReplay, self).__init__(message) class ALFREDImageDataset(IterableDataset): def __init__(self, args): # trajectories are scattered among several files # we gather their file names here self.trajectories = [] self.num_images = 0 self.image_size = args.image_width self.rotation_degrees = 90 # we have already the current frame self.rotation_steps = (360 // self.rotation_degrees) - 1 visited = set() with open(args.splits) as in_file: splits = json.load(in_file) stop = False for k, d in splits.items(): if args.split_id in k: for task in d: # load json file json_path = os.path.join(args.data_path, k, task['task'], 'traj_data.json') if json_path not in visited: visited.add(json_path) with open(json_path) as f: ex = json.load(f) # copy trajectory r_idx = task['repeat_idx'] # repeat_idx is the index of the annotation for each trajectory traj = ex.copy() # root & split traj['root'] = os.path.join(args.data_path, task['task']) traj['split'] = k traj['repeat_idx'] = r_idx traj['path'] = json_path self.trajectories.append(traj) self.num_images += len(traj['plan']['low_actions']) * (self.rotation_steps + 1) if args.debug and len(self.trajectories) == 3: stop = True break if stop: print("Debugging mode on!") break print(f"Discovered {len(self.trajectories)} files from the directory {args.data_path}") if args.debug: for traj in self.trajectories: print(traj["path"]) self.start = 0 self.end = len(self.trajectories) # assumes that we have a squared video frame self.image_loader = CustomImageLoader(min_size=args.image_width, max_size=args.image_width) self.failed = [] def __len__(self): return self.num_images def __iter__(self): worker_info = torch.utils.data.get_worker_info() # single worker if worker_info is None: # in this case we just process the entire dataset iter_start = self.start iter_end = self.end else: # we need to split the load among the workers making sure that there are no duplicates per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers))) worker_id = worker_info.id iter_start = self.start + worker_id * per_worker iter_end = min(iter_start + per_worker, self.end) # start THOR env env = ThorEnv(player_screen_width=args.image_width, player_screen_height=args.image_height) try: for idx in range(iter_start, iter_end): try: for root_dir, step_index, images, camera_infos, is_traj_finished in \ self._get_instance(env, self.trajectories[idx]): for pano_idx, (image, camera, done) in enumerate(zip(images, camera_infos, is_traj_finished)): norm_image, size = self.image_loader(image) yield root_dir, step_index, pano_idx, norm_image, size, camera, done except FailedReplay: self.failed.append(self.trajectories[idx]) except KeyboardInterrupt: pass finally: env.stop() def _get_instance(self, env, traj_data): json_file = traj_data["path"] # make directories root_dir = json_file.replace(TRAJ_DATA_JSON_FILENAME, "") # scene setup scene_num = traj_data['scene']['scene_num'] object_poses = traj_data['scene']['object_poses'] object_toggles = traj_data['scene']['object_toggles'] dirty_and_empty = traj_data['scene']['dirty_and_empty'] # reset scene_name = 'FloorPlan%d' % scene_num env.reset(scene_name) env.restore_scene(object_poses, object_toggles, dirty_and_empty) env.step(dict(traj_data['scene']['init_action'])) # setup task env.set_task(traj_data, args, reward_type='dense') num_actions = len(traj_data['plan']['low_actions']) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - scripts/generate_maskrcnn_horizon0.py [432:559]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - return panorama_frames, camera_info class FailedReplay(Exception): def __init__(self, message): super(FailedReplay, self).__init__(message) class ALFREDImageDataset(IterableDataset): def __init__(self, args): # trajectories are scattered among several files # we gather their file names here self.trajectories = [] self.num_images = 0 self.image_size = args.image_width self.rotation_degrees = 90 # we have already the current frame self.rotation_steps = (360 // self.rotation_degrees) - 1 visited = set() with open(args.splits) as in_file: splits = json.load(in_file) stop = False for k, d in splits.items(): if args.split_id in k: for task in d: # load json file json_path = os.path.join(args.data_path, k, task['task'], 'traj_data.json') if json_path not in visited: visited.add(json_path) with open(json_path) as f: ex = json.load(f) # copy trajectory r_idx = task['repeat_idx'] # repeat_idx is the index of the annotation for each trajectory traj = ex.copy() # root & split traj['root'] = os.path.join(args.data_path, task['task']) traj['split'] = k traj['repeat_idx'] = r_idx traj['path'] = json_path self.trajectories.append(traj) self.num_images += len(traj['plan']['low_actions']) * (self.rotation_steps + 1) if args.debug and len(self.trajectories) == 3: stop = True break if stop: print("Debugging mode on!") break print(f"Discovered {len(self.trajectories)} files from the directory {args.data_path}") if args.debug: for traj in self.trajectories: print(traj["path"]) self.start = 0 self.end = len(self.trajectories) # assumes that we have a squared video frame self.image_loader = CustomImageLoader(min_size=args.image_width, max_size=args.image_width) self.failed = [] def __len__(self): return self.num_images def __iter__(self): worker_info = torch.utils.data.get_worker_info() # single worker if worker_info is None: # in this case we just process the entire dataset iter_start = self.start iter_end = self.end else: # we need to split the load among the workers making sure that there are no duplicates per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers))) worker_id = worker_info.id iter_start = self.start + worker_id * per_worker iter_end = min(iter_start + per_worker, self.end) # start THOR env env = ThorEnv(player_screen_width=args.image_width, player_screen_height=args.image_height) try: for idx in range(iter_start, iter_end): try: for root_dir, step_index, images, camera_infos, is_traj_finished in \ self._get_instance(env, self.trajectories[idx]): for pano_idx, (image, camera, done) in enumerate(zip(images, camera_infos, is_traj_finished)): norm_image, size = self.image_loader(image) yield root_dir, step_index, pano_idx, norm_image, size, camera, done except FailedReplay: self.failed.append(self.trajectories[idx]) except KeyboardInterrupt: pass finally: env.stop() def _get_instance(self, env, traj_data): json_file = traj_data["path"] # make directories root_dir = json_file.replace(TRAJ_DATA_JSON_FILENAME, "") # scene setup scene_num = traj_data['scene']['scene_num'] object_poses = traj_data['scene']['object_poses'] object_toggles = traj_data['scene']['object_toggles'] dirty_and_empty = traj_data['scene']['dirty_and_empty'] # reset scene_name = 'FloorPlan%d' % scene_num env.reset(scene_name) env.restore_scene(object_poses, object_toggles, dirty_and_empty) env.step(dict(traj_data['scene']['init_action'])) # setup task env.set_task(traj_data, args, reward_type='dense') num_actions = len(traj_data['plan']['low_actions']) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -