vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/duster/dust3r/datasets/co3d.py

@@ -0,0 +1,165 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Dataloader for preprocessed Co3d_v2
+# dataset at https://github.com/facebookresearch/co3d - Creative Commons Attribution-NonCommercial 4.0 International
+# See datasets_preprocess/preprocess_co3d.py
+# --------------------------------------------------------
+import os.path as osp
+import json
+import itertools
+from collections import deque
+
+import cv2
+import numpy as np
+
+from dust3r.datasets.base.base_stereo_view_dataset import BaseStereoViewDataset
+from dust3r.utils.image import imread_cv2
+
+
+class Co3d(BaseStereoViewDataset):
+    def __init__(self, mask_bg=True, *args, ROOT, **kwargs):
+        self.ROOT = ROOT
+        super().__init__(*args, **kwargs)
+        assert mask_bg in (True, False, 'rand')
+        self.mask_bg = mask_bg
+        self.dataset_label = 'Co3d_v2'
+
+        # load all scenes
+        with open(osp.join(self.ROOT, f'selected_seqs_{self.split}.json'), 'r') as f:
+            self.scenes = json.load(f)
+            self.scenes = {k: v for k, v in self.scenes.items() if len(v) > 0}
+            self.scenes = {(k, k2): v2 for k, v in self.scenes.items()
+                           for k2, v2 in v.items()}
+        self.scene_list = list(self.scenes.keys())
+
+        # for each scene, we have 100 images ==> 360 degrees (so 25 frames ~= 90 degrees)
+        # we prepare all combinations such that i-j = +/- [5, 10, .., 90] degrees
+        self.combinations = [(i, j)
+                             for i, j in itertools.combinations(range(100), 2)
+                             if 0 < abs(i - j) <= 30 and abs(i - j) % 5 == 0]
+
+        self.invalidate = {scene: {} for scene in self.scene_list}
+
+    def __len__(self):
+        return len(self.scene_list) * len(self.combinations)
+
+    def _get_metadatapath(self, obj, instance, view_idx):
+        return osp.join(self.ROOT, obj, instance, 'images', f'frame{view_idx:06n}.npz')
+
+    def _get_impath(self, obj, instance, view_idx):
+        return osp.join(self.ROOT, obj, instance, 'images', f'frame{view_idx:06n}.jpg')
+
+    def _get_depthpath(self, obj, instance, view_idx):
+        return osp.join(self.ROOT, obj, instance, 'depths', f'frame{view_idx:06n}.jpg.geometric.png')
+
+    def _get_maskpath(self, obj, instance, view_idx):
+        return osp.join(self.ROOT, obj, instance, 'masks', f'frame{view_idx:06n}.png')
+
+    def _read_depthmap(self, depthpath, input_metadata):
+        depthmap = imread_cv2(depthpath, cv2.IMREAD_UNCHANGED)
+        depthmap = (depthmap.astype(np.float32) / 65535) * np.nan_to_num(input_metadata['maximum_depth'])
+        return depthmap
+
+    def _get_views(self, idx, resolution, rng):
+        # choose a scene
+        obj, instance = self.scene_list[idx // len(self.combinations)]
+        image_pool = self.scenes[obj, instance]
+        im1_idx, im2_idx = self.combinations[idx % len(self.combinations)]
+
+        # add a bit of randomness
+        last = len(image_pool) - 1
+
+        if resolution not in self.invalidate[obj, instance]:  # flag invalid images
+            self.invalidate[obj, instance][resolution] = [False for _ in range(len(image_pool))]
+
+        # decide now if we mask the bg
+        mask_bg = (self.mask_bg == True) or (self.mask_bg == 'rand' and rng.choice(2))
+
+        views = []
+        imgs_idxs = [max(0, min(im_idx + rng.integers(-4, 5), last)) for im_idx in [im2_idx, im1_idx]]
+        imgs_idxs = deque(imgs_idxs)
+        while len(imgs_idxs) > 0:  # some images (few) have zero depth
+            im_idx = imgs_idxs.pop()
+
+            if self.invalidate[obj, instance][resolution][im_idx]:
+                # search for a valid image
+                random_direction = 2 * rng.choice(2) - 1
+                for offset in range(1, len(image_pool)):
+                    tentative_im_idx = (im_idx + (random_direction * offset)) % len(image_pool)
+                    if not self.invalidate[obj, instance][resolution][tentative_im_idx]:
+                        im_idx = tentative_im_idx
+                        break
+
+            view_idx = image_pool[im_idx]
+
+            impath = self._get_impath(obj, instance, view_idx)
+            depthpath = self._get_depthpath(obj, instance, view_idx)
+
+            # load camera params
+            metadata_path = self._get_metadatapath(obj, instance, view_idx)
+            input_metadata = np.load(metadata_path)
+            camera_pose = input_metadata['camera_pose'].astype(np.float32)
+            intrinsics = input_metadata['camera_intrinsics'].astype(np.float32)
+
+            # load image and depth
+            rgb_image = imread_cv2(impath)
+            depthmap = self._read_depthmap(depthpath, input_metadata)
+
+            if mask_bg:
+                # load object mask
+                maskpath = self._get_maskpath(obj, instance, view_idx)
+                maskmap = imread_cv2(maskpath, cv2.IMREAD_UNCHANGED).astype(np.float32)
+                maskmap = (maskmap / 255.0) > 0.1
+
+                # update the depthmap with mask
+                depthmap *= maskmap
+
+            rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary(
+                rgb_image, depthmap, intrinsics, resolution, rng=rng, info=impath)
+
+            num_valid = (depthmap > 0.0).sum()
+            if num_valid == 0:
+                # problem, invalidate image and retry
+                self.invalidate[obj, instance][resolution][im_idx] = True
+                imgs_idxs.append(im_idx)
+                continue
+
+            views.append(dict(
+                img=rgb_image,
+                depthmap=depthmap,
+                camera_pose=camera_pose,
+                camera_intrinsics=intrinsics,
+                dataset=self.dataset_label,
+                label=osp.join(obj, instance),
+                instance=osp.split(impath)[1],
+            ))
+        return views
+
+
+if __name__ == "__main__":
+    from dust3r.datasets.base.base_stereo_view_dataset import view_name
+    from dust3r.viz import SceneViz, auto_cam_size
+    from dust3r.utils.image import rgb
+
+    dataset = Co3d(split='train', ROOT="data/co3d_subset_processed", resolution=224, aug_crop=16)
+
+    for idx in np.random.permutation(len(dataset)):
+        views = dataset[idx]
+        assert len(views) == 2
+        print(view_name(views[0]), view_name(views[1]))
+        viz = SceneViz()
+        poses = [views[view_idx]['camera_pose'] for view_idx in [0, 1]]
+        cam_size = max(auto_cam_size(poses), 0.001)
+        for view_idx in [0, 1]:
+            pts3d = views[view_idx]['pts3d']
+            valid_mask = views[view_idx]['valid_mask']
+            colors = rgb(views[view_idx]['img'])
+            viz.add_pointcloud(pts3d, colors, valid_mask)
+            viz.add_camera(pose_c2w=views[view_idx]['camera_pose'],
+                           focal=views[view_idx]['camera_intrinsics'][0, 0],
+                           color=(idx * 255, (1 - idx) * 255, 0),
+                           image=colors,
+                           cam_size=cam_size)
+        viz.show()

vismatch/third_party/duster/dust3r/datasets/habitat.py

@@ -0,0 +1,107 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Dataloader for preprocessed habitat
+# dataset at https://github.com/facebookresearch/habitat-sim/blob/main/DATASETS.md
+# See datasets_preprocess/habitat for more details
+# --------------------------------------------------------
+import os.path as osp
+import os
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"  # noqa
+import cv2  # noqa
+import numpy as np
+from PIL import Image
+import json
+
+from dust3r.datasets.base.base_stereo_view_dataset import BaseStereoViewDataset
+
+
+class Habitat(BaseStereoViewDataset):
+    def __init__(self, size, *args, ROOT, **kwargs):
+        self.ROOT = ROOT
+        super().__init__(*args, **kwargs)
+        assert self.split is not None
+        # loading list of scenes
+        with open(osp.join(self.ROOT, f'Habitat_{size}_scenes_{self.split}.txt')) as f:
+            self.scenes = f.read().splitlines()
+        self.instances = list(range(1, 5))
+
+    def filter_scene(self, label, instance=None):
+        if instance:
+            subscene, instance = instance.split('_')
+            label += '/' + subscene
+            self.instances = [int(instance) - 1]
+        valid = np.bool_([scene.startswith(label) for scene in self.scenes])
+        assert sum(valid), 'no scene was selected for {label=} {instance=}'
+        self.scenes = [scene for i, scene in enumerate(self.scenes) if valid[i]]
+
+    def _get_views(self, idx, resolution, rng):
+        scene = self.scenes[idx]
+        data_path, key = osp.split(osp.join(self.ROOT, scene))
+        views = []
+        two_random_views = [0, rng.choice(self.instances)]  # view 0 is connected with all other views
+        for view_index in two_random_views:
+            # load the view (and use the next one if this one's broken)
+            for ii in range(view_index, view_index + 5):
+                image, depthmap, intrinsics, camera_pose = self._load_one_view(data_path, key, ii % 5, resolution, rng)
+                if np.isfinite(camera_pose).all():
+                    break
+            views.append(dict(
+                img=image,
+                depthmap=depthmap,
+                camera_pose=camera_pose,  # cam2world
+                camera_intrinsics=intrinsics,
+                dataset='Habitat',
+                label=osp.relpath(data_path, self.ROOT),
+                instance=f"{key}_{view_index}"))
+        return views
+
+    def _load_one_view(self, data_path, key, view_index, resolution, rng):
+        view_index += 1  # file indices starts at 1
+        impath = osp.join(data_path, f"{key}_{view_index}.jpeg")
+        image = Image.open(impath)
+
+        depthmap_filename = osp.join(data_path, f"{key}_{view_index}_depth.exr")
+        depthmap = cv2.imread(depthmap_filename, cv2.IMREAD_GRAYSCALE | cv2.IMREAD_ANYDEPTH)
+
+        camera_params_filename = osp.join(data_path, f"{key}_{view_index}_camera_params.json")
+        with open(camera_params_filename, 'r') as f:
+            camera_params = json.load(f)
+
+        intrinsics = np.float32(camera_params['camera_intrinsics'])
+        camera_pose = np.eye(4, dtype=np.float32)
+        camera_pose[:3, :3] = camera_params['R_cam2world']
+        camera_pose[:3, 3] = camera_params['t_cam2world']
+
+        image, depthmap, intrinsics = self._crop_resize_if_necessary(
+            image, depthmap, intrinsics, resolution, rng, info=impath)
+        return image, depthmap, intrinsics, camera_pose
+
+
+if __name__ == "__main__":
+    from dust3r.datasets.base.base_stereo_view_dataset import view_name
+    from dust3r.viz import SceneViz, auto_cam_size
+    from dust3r.utils.image import rgb
+
+    dataset = Habitat(1_000_000, split='train', ROOT="data/habitat_processed",
+                      resolution=224, aug_crop=16)
+
+    for idx in np.random.permutation(len(dataset)):
+        views = dataset[idx]
+        assert len(views) == 2
+        print(view_name(views[0]), view_name(views[1]))
+        viz = SceneViz()
+        poses = [views[view_idx]['camera_pose'] for view_idx in [0, 1]]
+        cam_size = max(auto_cam_size(poses), 0.001)
+        for view_idx in [0, 1]:
+            pts3d = views[view_idx]['pts3d']
+            valid_mask = views[view_idx]['valid_mask']
+            colors = rgb(views[view_idx]['img'])
+            viz.add_pointcloud(pts3d, colors, valid_mask)
+            viz.add_camera(pose_c2w=views[view_idx]['camera_pose'],
+                           focal=views[view_idx]['camera_intrinsics'][0, 0],
+                           color=(idx * 255, (1 - idx) * 255, 0),
+                           image=colors,
+                           cam_size=cam_size)
+        viz.show()

vismatch/third_party/duster/dust3r/datasets/megadepth.py

@@ -0,0 +1,123 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Dataloader for preprocessed MegaDepth
+# dataset at https://www.cs.cornell.edu/projects/megadepth/
+# See datasets_preprocess/preprocess_megadepth.py
+# --------------------------------------------------------
+import os.path as osp
+import numpy as np
+
+from dust3r.datasets.base.base_stereo_view_dataset import BaseStereoViewDataset
+from dust3r.utils.image import imread_cv2
+
+
+class MegaDepth(BaseStereoViewDataset):
+    def __init__(self, *args, split, ROOT, **kwargs):
+        self.ROOT = ROOT
+        super().__init__(*args, **kwargs)
+        self.loaded_data = self._load_data(self.split)
+
+        if self.split is None:
+            pass
+        elif self.split == 'train':
+            self.select_scene(('0015', '0022'), opposite=True)
+        elif self.split == 'val':
+            self.select_scene(('0015', '0022'))
+        else:
+            raise ValueError(f'bad {self.split=}')
+
+    def _load_data(self, split):
+        with np.load(osp.join(self.ROOT, 'all_metadata.npz')) as data:
+            self.all_scenes = data['scenes']
+            self.all_images = data['images']
+            self.pairs = data['pairs']
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def get_stats(self):
+        return f'{len(self)} pairs from {len(self.all_scenes)} scenes'
+
+    def select_scene(self, scene, *instances, opposite=False):
+        scenes = (scene,) if isinstance(scene, str) else tuple(scene)
+        scene_id = [s.startswith(scenes) for s in self.all_scenes]
+        assert any(scene_id), 'no scene found'
+
+        valid = np.in1d(self.pairs['scene_id'], np.nonzero(scene_id)[0])
+        if instances:
+            image_id = [i.startswith(instances) for i in self.all_images]
+            image_id = np.nonzero(image_id)[0]
+            assert len(image_id), 'no instance found'
+            # both together?
+            if len(instances) == 2:
+                valid &= np.in1d(self.pairs['im1_id'], image_id) & np.in1d(self.pairs['im2_id'], image_id)
+            else:
+                valid &= np.in1d(self.pairs['im1_id'], image_id) | np.in1d(self.pairs['im2_id'], image_id)
+
+        if opposite:
+            valid = ~valid
+        assert valid.any()
+        self.pairs = self.pairs[valid]
+
+    def _get_views(self, pair_idx, resolution, rng):
+        scene_id, im1_id, im2_id, score = self.pairs[pair_idx]
+
+        scene, subscene = self.all_scenes[scene_id].split()
+        seq_path = osp.join(self.ROOT, scene, subscene)
+
+        views = []
+
+        for im_id in [im1_id, im2_id]:
+            img = self.all_images[im_id]
+            try:
+                image = imread_cv2(osp.join(seq_path, img + '.jpg'))
+                depthmap = imread_cv2(osp.join(seq_path, img + ".exr"))
+                camera_params = np.load(osp.join(seq_path, img + ".npz"))
+            except Exception as e:
+                raise OSError(f'cannot load {img}, got exception {e}')
+
+            intrinsics = np.float32(camera_params['intrinsics'])
+            camera_pose = np.float32(camera_params['cam2world'])
+
+            image, depthmap, intrinsics = self._crop_resize_if_necessary(
+                image, depthmap, intrinsics, resolution, rng, info=(seq_path, img))
+
+            views.append(dict(
+                img=image,
+                depthmap=depthmap,
+                camera_pose=camera_pose,  # cam2world
+                camera_intrinsics=intrinsics,
+                dataset='MegaDepth',
+                label=osp.relpath(seq_path, self.ROOT),
+                instance=img))
+
+        return views
+
+
+if __name__ == "__main__":
+    from dust3r.datasets.base.base_stereo_view_dataset import view_name
+    from dust3r.viz import SceneViz, auto_cam_size
+    from dust3r.utils.image import rgb
+
+    dataset = MegaDepth(split='train', ROOT="data/megadepth_processed", resolution=224, aug_crop=16)
+
+    for idx in np.random.permutation(len(dataset)):
+        views = dataset[idx]
+        assert len(views) == 2
+        print(idx, view_name(views[0]), view_name(views[1]))
+        viz = SceneViz()
+        poses = [views[view_idx]['camera_pose'] for view_idx in [0, 1]]
+        cam_size = max(auto_cam_size(poses), 0.001)
+        for view_idx in [0, 1]:
+            pts3d = views[view_idx]['pts3d']
+            valid_mask = views[view_idx]['valid_mask']
+            colors = rgb(views[view_idx]['img'])
+            viz.add_pointcloud(pts3d, colors, valid_mask)
+            viz.add_camera(pose_c2w=views[view_idx]['camera_pose'],
+                           focal=views[view_idx]['camera_intrinsics'][0, 0],
+                           color=(idx * 255, (1 - idx) * 255, 0),
+                           image=colors,
+                           cam_size=cam_size)
+        viz.show()

vismatch/third_party/duster/dust3r/datasets/scannetpp.py

@@ -0,0 +1,96 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Dataloader for preprocessed scannet++
+# dataset at https://github.com/scannetpp/scannetpp - non-commercial research and educational purposes
+# https://kaldir.vc.in.tum.de/scannetpp/static/scannetpp-terms-of-use.pdf
+# See datasets_preprocess/preprocess_scannetpp.py
+# --------------------------------------------------------
+import os.path as osp
+import cv2
+import numpy as np
+
+from dust3r.datasets.base.base_stereo_view_dataset import BaseStereoViewDataset
+from dust3r.utils.image import imread_cv2
+
+
+class ScanNetpp(BaseStereoViewDataset):
+    def __init__(self, *args, ROOT, **kwargs):
+        self.ROOT = ROOT
+        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']
+            self.sceneids = data['sceneids']
+            self.images = data['images']
+            self.intrinsics = data['intrinsics'].astype(np.float32)
+            self.trajectories = data['trajectories'].astype(np.float32)
+            self.pairs = data['pairs'][:, :2].astype(int)
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def _get_views(self, idx, resolution, rng):
+
+        image_idx1, image_idx2 = self.pairs[idx]
+
+        views = []
+        for view_idx in [image_idx1, image_idx2]:
+            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)
+
+            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)}',
+            ))
+        return views
+
+
+if __name__ == "__main__":
+    from dust3r.datasets.base.base_stereo_view_dataset import view_name
+    from dust3r.viz import SceneViz, auto_cam_size
+    from dust3r.utils.image import rgb
+
+    dataset = ScanNetpp(split='train', ROOT="data/scannetpp_processed", resolution=224, aug_crop=16)
+
+    for idx in np.random.permutation(len(dataset)):
+        views = dataset[idx]
+        assert len(views) == 2
+        print(view_name(views[0]), view_name(views[1]))
+        viz = SceneViz()
+        poses = [views[view_idx]['camera_pose'] for view_idx in [0, 1]]
+        cam_size = max(auto_cam_size(poses), 0.001)
+        for view_idx in [0, 1]:
+            pts3d = views[view_idx]['pts3d']
+            valid_mask = views[view_idx]['valid_mask']
+            colors = rgb(views[view_idx]['img'])
+            viz.add_pointcloud(pts3d, colors, valid_mask)
+            viz.add_camera(pose_c2w=views[view_idx]['camera_pose'],
+                           focal=views[view_idx]['camera_intrinsics'][0, 0],
+                           color=(idx*255, (1 - idx)*255, 0),
+                           image=colors,
+                           cam_size=cam_size)
+        viz.show()

vismatch/third_party/duster/dust3r/datasets/staticthings3d.py

@@ -0,0 +1,96 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Dataloader for preprocessed StaticThings3D
+# dataset at https://github.com/lmb-freiburg/robustmvd/
+# See datasets_preprocess/preprocess_staticthings3d.py
+# --------------------------------------------------------
+import os.path as osp
+import numpy as np
+
+from dust3r.datasets.base.base_stereo_view_dataset import BaseStereoViewDataset
+from dust3r.utils.image import imread_cv2
+
+
+class StaticThings3D (BaseStereoViewDataset):
+    """ Dataset of indoor scenes, 5 images each time
+    """
+    def __init__(self, ROOT, *args, mask_bg='rand', **kwargs):
+        self.ROOT = ROOT
+        super().__init__(*args, **kwargs)
+
+        assert mask_bg in (True, False, 'rand')
+        self.mask_bg = mask_bg
+
+        # loading all pairs
+        assert self.split is None
+        self.pairs = np.load(osp.join(ROOT, 'staticthings_pairs.npy'))
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def get_stats(self):
+        return f'{len(self)} pairs'
+
+    def _get_views(self, pair_idx, resolution, rng):
+        scene, seq, cam1, im1, cam2, im2 = self.pairs[pair_idx]
+        seq_path = osp.join('TRAIN', scene.decode('ascii'), f'{seq:04d}')
+
+        views = []
+
+        mask_bg = (self.mask_bg == True) or (self.mask_bg == 'rand' and rng.choice(2))
+
+        CAM = {b'l':'left', b'r':'right'}
+        for cam, idx in [(CAM[cam1], im1), (CAM[cam2], im2)]:
+            num = f"{idx:04n}"
+            img = num+"_clean.jpg" if rng.choice(2) else num+"_final.jpg"
+            image = imread_cv2(osp.join(self.ROOT, seq_path, cam, img))
+            depthmap = imread_cv2(osp.join(self.ROOT, seq_path, cam, num+".exr"))
+            camera_params = np.load(osp.join(self.ROOT, seq_path, cam, num+".npz"))
+
+            intrinsics = camera_params['intrinsics']
+            camera_pose = camera_params['cam2world']
+
+            if mask_bg:
+                depthmap[depthmap > 200] = 0
+
+            image, depthmap, intrinsics = self._crop_resize_if_necessary(image, depthmap, intrinsics, resolution, rng, info=(seq_path,cam,img))
+
+            views.append(dict(
+                img = image, 
+                depthmap = depthmap,
+                camera_pose = camera_pose, # cam2world
+                camera_intrinsics = intrinsics,
+                dataset = 'StaticThings3D',
+                label = seq_path,
+                instance = cam+'_'+img))
+
+        return views
+
+
+if __name__ == '__main__':
+    from dust3r.datasets.base.base_stereo_view_dataset import view_name
+    from dust3r.viz import SceneViz, auto_cam_size
+    from dust3r.utils.image import rgb
+
+    dataset = StaticThings3D(ROOT="data/staticthings3d_processed", resolution=224, aug_crop=16)
+
+    for idx in np.random.permutation(len(dataset)):
+        views = dataset[idx]
+        assert len(views) == 2
+        print(idx, view_name(views[0]), view_name(views[1]))
+        viz = SceneViz()
+        poses = [views[view_idx]['camera_pose'] for view_idx in [0, 1]]
+        cam_size = max(auto_cam_size(poses), 0.001)
+        for view_idx in [0, 1]:
+            pts3d = views[view_idx]['pts3d']
+            valid_mask = views[view_idx]['valid_mask']
+            colors = rgb(views[view_idx]['img'])
+            viz.add_pointcloud(pts3d, colors, valid_mask)
+            viz.add_camera(pose_c2w=views[view_idx]['camera_pose'],
+                           focal=views[view_idx]['camera_intrinsics'][0, 0],
+                           color=(idx*255, (1 - idx)*255, 0),
+                           image=colors,
+                           cam_size=cam_size)
+        viz.show()

vismatch/third_party/duster/dust3r/datasets/utils/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).