Spaces:
Runtime error
Runtime error
| import os.path as osp | |
| import cv2 | |
| import numpy as np | |
| import itertools | |
| import os | |
| import sys | |
| sys.path.append(osp.join(osp.dirname(__file__), "..", "..")) | |
| from dust3r.datasets.base.base_multiview_dataset import BaseMultiViewDataset | |
| from dust3r.utils.image import imread_cv2 | |
| class ScanNetpp_Multi(BaseMultiViewDataset): | |
| def __init__(self, *args, ROOT, **kwargs): | |
| self.ROOT = ROOT | |
| self.video = True | |
| self.is_metric = True | |
| self.max_interval = 3 | |
| super().__init__(*args, **kwargs) | |
| assert self.split == "train" | |
| self.loaded_data = self._load_data() | |
| def _load_data(self): | |
| with np.load(osp.join(self.ROOT, "all_metadata.npz")) as data: | |
| self.scenes = data["scenes"] | |
| offset = 0 | |
| scenes = [] | |
| sceneids = [] | |
| images = [] | |
| intrinsics = [] | |
| trajectories = [] | |
| groups = [] | |
| id_ranges = [] | |
| j = 0 | |
| self.image_num = 0 | |
| for scene in self.scenes: | |
| scene_dir = osp.join(self.ROOT, scene) | |
| with np.load( | |
| osp.join(scene_dir, "new_scene_metadata.npz"), allow_pickle=True | |
| ) as data: | |
| imgs = data["images"] | |
| self.image_num += len(imgs) | |
| img_ids = np.arange(len(imgs)).tolist() | |
| intrins = data["intrinsics"] | |
| traj = data["trajectories"] | |
| imgs_on_disk = sorted(os.listdir(osp.join(scene_dir, "images"))) | |
| imgs_on_disk = list(map(lambda x: x[:-4], imgs_on_disk)) | |
| dslr_ids = [ | |
| i + offset | |
| for i in img_ids | |
| if imgs[i].startswith("DSC") and imgs[i] in imgs_on_disk | |
| ] | |
| iphone_ids = [ | |
| i + offset | |
| for i in img_ids | |
| if imgs[i].startswith("frame") and imgs[i] in imgs_on_disk | |
| ] | |
| num_imgs = len(imgs) | |
| assert max(dslr_ids) < min(iphone_ids) | |
| assert "image_collection" in data | |
| img_groups = [] | |
| img_id_ranges = [] | |
| for ref_id, group in data["image_collection"].item().items(): | |
| if len(group) + 1 < self.num_views: | |
| continue | |
| group.insert(0, (ref_id, 1.0)) | |
| sorted_group = sorted(group, key=lambda x: x[1], reverse=True) | |
| group = [int(x[0] + offset) for x in sorted_group] | |
| img_groups.append(sorted(group)) | |
| if imgs[ref_id].startswith("frame"): | |
| img_id_ranges.append(dslr_ids) | |
| else: | |
| img_id_ranges.append(iphone_ids) | |
| if len(img_groups) == 0: | |
| print(f"Skipping {scene}") | |
| continue | |
| scenes.append(scene) | |
| sceneids.extend([j] * num_imgs) | |
| images.extend(imgs) | |
| intrinsics.append(intrins) | |
| trajectories.append(traj) | |
| # offset groups | |
| groups.extend(img_groups) | |
| id_ranges.extend(img_id_ranges) | |
| offset += num_imgs | |
| j += 1 | |
| self.scenes = scenes | |
| self.sceneids = sceneids | |
| self.images = images | |
| self.intrinsics = np.concatenate(intrinsics, axis=0) | |
| self.trajectories = np.concatenate(trajectories, axis=0) | |
| self.id_ranges = id_ranges | |
| self.groups = groups | |
| def __len__(self): | |
| return len(self.groups) * 10 | |
| def get_image_num(self): | |
| return self.image_num | |
| def _get_views(self, idx, resolution, rng, num_views): | |
| idx = idx // 10 | |
| image_idxs = self.groups[idx] | |
| rand_val = rng.random() | |
| image_idxs_video = self.id_ranges[idx] | |
| cut_off = num_views if not self.allow_repeat else max(num_views // 3, 3) | |
| start_image_idxs = image_idxs_video[: len(image_idxs_video) - cut_off + 1] | |
| if rand_val < 0.7 and len(start_image_idxs) > 0: | |
| start_id = rng.choice(start_image_idxs) | |
| pos, ordered_video = self.get_seq_from_start_id( | |
| num_views, | |
| start_id, | |
| image_idxs_video, | |
| rng, | |
| max_interval=self.max_interval, | |
| video_prob=0.8, | |
| fix_interval_prob=0.5, | |
| block_shuffle=16, | |
| ) | |
| image_idxs = np.array(image_idxs_video)[pos] | |
| else: | |
| ordered_video = True | |
| # ordered video with varying intervals | |
| num_candidates = len(image_idxs) | |
| max_id = min(num_candidates, int(num_views * (2 + 2 * rng.random()))) | |
| image_idxs = sorted(rng.permutation(image_idxs[:max_id])[:num_views]) | |
| if rand_val > 0.75: | |
| ordered_video = False | |
| image_idxs = rng.permutation(image_idxs) | |
| views = [] | |
| for v, view_idx in enumerate(image_idxs): | |
| scene_id = self.sceneids[view_idx] | |
| scene_dir = osp.join(self.ROOT, self.scenes[scene_id]) | |
| intrinsics = self.intrinsics[view_idx] | |
| camera_pose = self.trajectories[view_idx] | |
| basename = self.images[view_idx] | |
| # Load RGB image | |
| rgb_image = imread_cv2(osp.join(scene_dir, "images", basename + ".jpg")) | |
| # Load depthmap | |
| depthmap = imread_cv2( | |
| osp.join(scene_dir, "depth", basename + ".png"), cv2.IMREAD_UNCHANGED | |
| ) | |
| depthmap = depthmap.astype(np.float32) / 1000 | |
| depthmap[~np.isfinite(depthmap)] = 0 # invalid | |
| rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary( | |
| rgb_image, depthmap, intrinsics, resolution, rng=rng, info=view_idx | |
| ) | |
| # generate img mask and raymap mask | |
| img_mask, ray_mask = self.get_img_and_ray_masks( | |
| self.is_metric, v, rng, p=[0.75, 0.2, 0.05] | |
| ) | |
| views.append( | |
| dict( | |
| img=rgb_image, | |
| depthmap=depthmap.astype(np.float32), | |
| camera_pose=camera_pose.astype(np.float32), | |
| camera_intrinsics=intrinsics.astype(np.float32), | |
| dataset="ScanNet++", | |
| label=self.scenes[scene_id] + "_" + basename, | |
| instance=f"{str(idx)}_{str(view_idx)}", | |
| is_metric=self.is_metric, | |
| is_video=ordered_video, | |
| quantile=np.array(0.99, dtype=np.float32), | |
| img_mask=img_mask, | |
| ray_mask=ray_mask, | |
| camera_only=False, | |
| depth_only=False, | |
| single_view=False, | |
| reset=False, | |
| ) | |
| ) | |
| assert len(views) == num_views | |
| return views | |