(category, wildrgbd_dir, output_dir, img_size, split, max_num_sequences_per_object,
output_num_frames, seed)
| 58 | |
| 59 | |
| 60 | def prepare_sequences(category, wildrgbd_dir, output_dir, img_size, split, max_num_sequences_per_object, |
| 61 | output_num_frames, seed): |
| 62 | random.seed(seed) |
| 63 | category_dir = osp.join(wildrgbd_dir, category) |
| 64 | category_output_dir = osp.join(output_dir, category) |
| 65 | sequences_all = get_set_list(category_dir, split) |
| 66 | sequences_all = sorted(sequences_all) |
| 67 | |
| 68 | sequences_all_tmp = [] |
| 69 | for seq_name in sequences_all: |
| 70 | scene_dir = osp.join(wildrgbd_dir, category_dir, seq_name) |
| 71 | if not os.path.isdir(scene_dir): |
| 72 | print(f'{scene_dir} does not exist, skipped') |
| 73 | continue |
| 74 | sequences_all_tmp.append(seq_name) |
| 75 | sequences_all = sequences_all_tmp |
| 76 | if len(sequences_all) <= max_num_sequences_per_object: |
| 77 | selected_sequences = sequences_all |
| 78 | else: |
| 79 | selected_sequences = random.sample(sequences_all, max_num_sequences_per_object) |
| 80 | |
| 81 | selected_sequences_numbers_dict = {} |
| 82 | for seq_name in tqdm(selected_sequences, leave=False): |
| 83 | scene_dir = osp.join(category_dir, seq_name) |
| 84 | scene_output_dir = osp.join(category_output_dir, seq_name) |
| 85 | with open(osp.join(scene_dir, 'metadata'), 'r') as f: |
| 86 | metadata = json.load(f) |
| 87 | |
| 88 | K = np.array(metadata["K"]).reshape(3, 3).T |
| 89 | fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2] |
| 90 | w, h = metadata["w"], metadata["h"] |
| 91 | |
| 92 | camera_intrinsics = np.array( |
| 93 | [[fx, 0, cx], |
| 94 | [0, fy, cy], |
| 95 | [0, 0, 1]] |
| 96 | ) |
| 97 | camera_to_world_path = os.path.join(scene_dir, 'cam_poses.txt') |
| 98 | camera_to_world_content = np.genfromtxt(camera_to_world_path) |
| 99 | camera_to_world = camera_to_world_content[:, 1:].reshape(-1, 4, 4) |
| 100 | |
| 101 | frame_idx = camera_to_world_content[:, 0] |
| 102 | num_frames = frame_idx.shape[0] |
| 103 | assert num_frames >= output_num_frames |
| 104 | assert np.all(frame_idx == np.arange(num_frames)) |
| 105 | |
| 106 | # selected_sequences_numbers_dict[seq_name] = num_frames |
| 107 | |
| 108 | selected_frames = np.round(np.linspace(0, num_frames - 1, output_num_frames)).astype(int).tolist() |
| 109 | selected_sequences_numbers_dict[seq_name] = selected_frames |
| 110 | |
| 111 | for frame_id in tqdm(selected_frames): |
| 112 | depth_path = os.path.join(scene_dir, 'depth', f'{frame_id:0>5d}.png') |
| 113 | masks_path = os.path.join(scene_dir, 'masks', f'{frame_id:0>5d}.png') |
| 114 | rgb_path = os.path.join(scene_dir, 'rgb', f'{frame_id:0>5d}.png') |
| 115 | |
| 116 | input_rgb_image = PIL.Image.open(rgb_path).convert('RGB') |
| 117 | input_mask = plt.imread(masks_path) |
no test coverage detected