| 19 | |
| 20 | |
| 21 | class Co3d(BaseStereoViewDataset): |
| 22 | def __init__(self, mask_bg=True, *args, ROOT, **kwargs): |
| 23 | self.ROOT = ROOT |
| 24 | super().__init__(*args, **kwargs) |
| 25 | assert mask_bg in (True, False, 'rand') |
| 26 | self.mask_bg = mask_bg |
| 27 | self.dataset_label = 'Co3d_v2' |
| 28 | |
| 29 | # load all scenes |
| 30 | with open(osp.join(self.ROOT, f'selected_seqs_{self.split}.json'), 'r') as f: |
| 31 | self.scenes = json.load(f) |
| 32 | self.scenes = {k: v for k, v in self.scenes.items() if len(v) > 0} |
| 33 | self.scenes = {(k, k2): v2 for k, v in self.scenes.items() |
| 34 | for k2, v2 in v.items()} |
| 35 | self.scene_list = list(self.scenes.keys()) |
| 36 | |
| 37 | # for each scene, we have 100 images ==> 360 degrees (so 25 frames ~= 90 degrees) |
| 38 | # we prepare all combinations such that i-j = +/- [5, 10, .., 90] degrees |
| 39 | self.combinations = [(i, j) |
| 40 | for i, j in itertools.combinations(range(100), 2) |
| 41 | if 0 < abs(i - j) <= 30 and abs(i - j) % 5 == 0] |
| 42 | |
| 43 | self.invalidate = {scene: {} for scene in self.scene_list} |
| 44 | |
| 45 | def __len__(self): |
| 46 | return len(self.scene_list) * len(self.combinations) |
| 47 | |
| 48 | def _get_metadatapath(self, obj, instance, view_idx): |
| 49 | return osp.join(self.ROOT, obj, instance, 'images', f'frame{view_idx:06n}.npz') |
| 50 | |
| 51 | def _get_impath(self, obj, instance, view_idx): |
| 52 | return osp.join(self.ROOT, obj, instance, 'images', f'frame{view_idx:06n}.jpg') |
| 53 | |
| 54 | def _get_depthpath(self, obj, instance, view_idx): |
| 55 | return osp.join(self.ROOT, obj, instance, 'depths', f'frame{view_idx:06n}.jpg.geometric.png') |
| 56 | |
| 57 | def _get_maskpath(self, obj, instance, view_idx): |
| 58 | return osp.join(self.ROOT, obj, instance, 'masks', f'frame{view_idx:06n}.png') |
| 59 | |
| 60 | def _read_depthmap(self, depthpath, input_metadata): |
| 61 | depthmap = imread_cv2(depthpath, cv2.IMREAD_UNCHANGED) |
| 62 | depthmap = (depthmap.astype(np.float32) / 65535) * np.nan_to_num(input_metadata['maximum_depth']) |
| 63 | return depthmap |
| 64 | |
| 65 | def _get_views(self, idx, resolution, rng): |
| 66 | # choose a scene |
| 67 | obj, instance = self.scene_list[idx // len(self.combinations)] |
| 68 | image_pool = self.scenes[obj, instance] |
| 69 | im1_idx, im2_idx = self.combinations[idx % len(self.combinations)] |
| 70 | |
| 71 | # add a bit of randomness |
| 72 | last = len(image_pool) - 1 |
| 73 | |
| 74 | if resolution not in self.invalidate[obj, instance]: # flag invalid images |
| 75 | self.invalidate[obj, instance][resolution] = [False for _ in range(len(image_pool))] |
| 76 | |
| 77 | # decide now if we mask the bg |
| 78 | mask_bg = (self.mask_bg == True) or (self.mask_bg == 'rand' and rng.choice(2)) |