vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/mast3r/dust3r/datasets_preprocess/habitat/habitat_renderer/habitat_sim_envmaps_renderer.py

@@ -0,0 +1,170 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Render environment maps from 3D meshes using the Habitat Sim simulator.
+# --------------------------------------------------------
+import numpy as np
+import habitat_sim
+import math
+from habitat_renderer import projections
+
+# OpenCV to habitat camera convention transformation
+R_OPENCV2HABITAT = np.stack((habitat_sim.geo.RIGHT, -habitat_sim.geo.UP, habitat_sim.geo.FRONT), axis=0)
+
+CUBEMAP_FACE_LABELS = ["left", "front", "right", "back", "up", "down"]
+# Expressed while considering Habitat coordinates systems
+CUBEMAP_FACE_ORIENTATIONS_ROTVEC = [
+    [0, math.pi / 2, 0],  # Left
+    [0, 0, 0],  # Front
+                [0, - math.pi / 2, 0],  # Right
+                [0, math.pi, 0],  # Back
+                [math.pi / 2, 0, 0],  # Up
+                [-math.pi / 2, 0, 0],]  # Down
+
+class NoNaviguableSpaceError(RuntimeError):
+    def __init__(self, *args):
+        super().__init__(*args)
+
+class HabitatEnvironmentMapRenderer:
+    def __init__(self,
+                 scene,
+                 navmesh,
+                 scene_dataset_config_file,
+                 render_equirectangular=False,
+                 equirectangular_resolution=(512, 1024),
+                 render_cubemap=False,
+                 cubemap_resolution=(512, 512),
+                 render_depth=False,
+                 gpu_id=0):
+        self.scene = scene
+        self.navmesh = navmesh
+        self.scene_dataset_config_file = scene_dataset_config_file
+        self.gpu_id = gpu_id
+
+        self.render_equirectangular = render_equirectangular
+        self.equirectangular_resolution = equirectangular_resolution
+        self.equirectangular_projection = projections.EquirectangularProjection(*equirectangular_resolution)
+        # 3D unit ray associated to each pixel of the equirectangular map
+        equirectangular_rays = projections.get_projection_rays(self.equirectangular_projection)
+        # Not needed, but just in case.
+        equirectangular_rays /= np.linalg.norm(equirectangular_rays, axis=-1, keepdims=True)
+        # Depth map created by Habitat are produced by warping a cubemap,
+        # so the values do not correspond to distance to the center and need some scaling.
+        self.equirectangular_depth_scale_factors = 1.0 / np.max(np.abs(equirectangular_rays), axis=-1)
+
+        self.render_cubemap = render_cubemap
+        self.cubemap_resolution = cubemap_resolution
+
+        self.render_depth = render_depth
+
+        self.seed = None
+        self._lazy_initialization()
+
+    def _lazy_initialization(self):
+        # Lazy random seeding and instantiation of the simulator to deal with multiprocessing properly
+        if self.seed == None:
+            # Re-seed numpy generator
+            np.random.seed()
+            self.seed = np.random.randint(2**32-1)
+            sim_cfg = habitat_sim.SimulatorConfiguration()
+            sim_cfg.scene_id = self.scene
+            if self.scene_dataset_config_file is not None and self.scene_dataset_config_file != "":
+                sim_cfg.scene_dataset_config_file = self.scene_dataset_config_file
+            sim_cfg.random_seed = self.seed
+            sim_cfg.load_semantic_mesh = False
+            sim_cfg.gpu_device_id = self.gpu_id
+
+            sensor_specifications = []
+
+            # Add cubemaps
+            if self.render_cubemap:
+                for face_id, orientation in enumerate(CUBEMAP_FACE_ORIENTATIONS_ROTVEC):
+                    rgb_sensor_spec = habitat_sim.CameraSensorSpec()
+                    rgb_sensor_spec.uuid = f"color_cubemap_{CUBEMAP_FACE_LABELS[face_id]}"
+                    rgb_sensor_spec.sensor_type = habitat_sim.SensorType.COLOR
+                    rgb_sensor_spec.resolution = self.cubemap_resolution
+                    rgb_sensor_spec.hfov = 90
+                    rgb_sensor_spec.position = [0.0, 0.0, 0.0]
+                    rgb_sensor_spec.orientation = orientation
+                    sensor_specifications.append(rgb_sensor_spec)
+
+                    if self.render_depth:
+                        depth_sensor_spec = habitat_sim.CameraSensorSpec()
+                        depth_sensor_spec.uuid = f"depth_cubemap_{CUBEMAP_FACE_LABELS[face_id]}"
+                        depth_sensor_spec.sensor_type = habitat_sim.SensorType.DEPTH
+                        depth_sensor_spec.resolution = self.cubemap_resolution
+                        depth_sensor_spec.hfov = 90
+                        depth_sensor_spec.position = [0.0, 0.0, 0.0]
+                        depth_sensor_spec.orientation = orientation
+                        sensor_specifications.append(depth_sensor_spec)
+
+            # Add equirectangular map
+            if self.render_equirectangular:
+                rgb_sensor_spec = habitat_sim.bindings.EquirectangularSensorSpec()
+                rgb_sensor_spec.uuid = "color_equirectangular"
+                rgb_sensor_spec.resolution = self.equirectangular_resolution
+                rgb_sensor_spec.position = [0.0, 0.0, 0.0]
+                sensor_specifications.append(rgb_sensor_spec)
+
+                if self.render_depth:
+                    depth_sensor_spec = habitat_sim.bindings.EquirectangularSensorSpec()
+                    depth_sensor_spec.uuid = "depth_equirectangular"
+                    depth_sensor_spec.sensor_type = habitat_sim.SensorType.DEPTH
+                    depth_sensor_spec.resolution = self.equirectangular_resolution
+                    depth_sensor_spec.position = [0.0, 0.0, 0.0]
+                    depth_sensor_spec.orientation
+                    sensor_specifications.append(depth_sensor_spec)
+
+            agent_cfg = habitat_sim.agent.AgentConfiguration(sensor_specifications=sensor_specifications)
+
+            cfg = habitat_sim.Configuration(sim_cfg, [agent_cfg])
+            self.sim = habitat_sim.Simulator(cfg)
+            if self.navmesh is not None and self.navmesh != "":
+                # Use pre-computed navmesh (the one generated automatically does some weird stuffs like going on top of the roof)
+                # See https://youtu.be/kunFMRJAu2U?t=1522 regarding navmeshes
+                self.sim.pathfinder.load_nav_mesh(self.navmesh)
+
+            # Check that the navmesh is not empty
+            if not self.sim.pathfinder.is_loaded:
+                # Try to compute a navmesh
+                navmesh_settings = habitat_sim.NavMeshSettings()
+                navmesh_settings.set_defaults()
+                self.sim.recompute_navmesh(self.sim.pathfinder, navmesh_settings, True)
+
+            # Check that the navmesh is not empty
+            if not self.sim.pathfinder.is_loaded:
+                raise NoNaviguableSpaceError(f"No naviguable location (scene: {self.scene} -- navmesh: {self.navmesh})")
+
+            self.agent = self.sim.initialize_agent(agent_id=0)
+
+    def close(self):
+        if hasattr(self, 'sim'):
+            self.sim.close()
+
+    def __del__(self):
+        self.close()
+
+    def render_viewpoint(self, viewpoint_position):
+        agent_state = habitat_sim.AgentState()
+        agent_state.position = viewpoint_position
+        # agent_state.rotation = viewpoint_orientation
+        self.agent.set_state(agent_state)
+        viewpoint_observations = self.sim.get_sensor_observations(agent_ids=0)
+
+        try:
+            # Depth map values have been obtained using cubemap rendering internally,
+            # so they do not really correspond to distance to the viewpoint in practice
+            # and they need some scaling
+            viewpoint_observations["depth_equirectangular"] *= self.equirectangular_depth_scale_factors
+        except KeyError:
+            pass
+
+        data = dict(observations=viewpoint_observations, position=viewpoint_position)
+        return data
+
+    def up_direction(self):
+        return np.asarray(habitat_sim.geo.UP).tolist()
+    
+    def R_cam_to_world(self):
+        return R_OPENCV2HABITAT.tolist()

vismatch/third_party/mast3r/dust3r/datasets_preprocess/habitat/habitat_renderer/multiview_crop_generator.py

@@ -0,0 +1,93 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Generate pairs of crops from a dataset of environment maps.
+# --------------------------------------------------------
+import os
+import numpy as np
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"  # noqa
+import cv2
+import collections
+from habitat_renderer import projections, projections_conversions
+from habitat_renderer.habitat_sim_envmaps_renderer import HabitatEnvironmentMapRenderer
+
+ViewpointData = collections.namedtuple("ViewpointData", ["colormap", "distancemap", "pointmap", "position"])
+
+class HabitatMultiviewCrops:
+    def __init__(self,
+                 scene,
+                 navmesh,
+                 scene_dataset_config_file,
+                 equirectangular_resolution=(400, 800),
+                 crop_resolution=(240, 320),
+                 pixel_jittering_iterations=5,
+                 jittering_noise_level=1.0):
+        self.crop_resolution = crop_resolution
+
+        self.pixel_jittering_iterations = pixel_jittering_iterations
+        self.jittering_noise_level = jittering_noise_level
+
+        # Instanciate the low resolution habitat sim renderer
+        self.lowres_envmap_renderer = HabitatEnvironmentMapRenderer(scene=scene,
+                                                                    navmesh=navmesh,
+                                                                    scene_dataset_config_file=scene_dataset_config_file,
+                                                                    equirectangular_resolution=equirectangular_resolution,
+                                                                    render_depth=True,
+                                                                    render_equirectangular=True)
+        self.R_cam_to_world = np.asarray(self.lowres_envmap_renderer.R_cam_to_world())
+        self.up_direction = np.asarray(self.lowres_envmap_renderer.up_direction())
+
+        # Projection applied by each environment map
+        self.envmap_height, self.envmap_width = self.lowres_envmap_renderer.equirectangular_resolution
+        base_projection = projections.EquirectangularProjection(self.envmap_height, self.envmap_width)
+        self.envmap_projection = projections.RotatedProjection(base_projection, self.R_cam_to_world.T)
+        # 3D Rays map associated to each envmap
+        self.envmap_rays = projections.get_projection_rays(self.envmap_projection)
+
+    def compute_pointmap(self, distancemap, position):
+        # Point cloud associated to each ray
+        return self.envmap_rays * distancemap[:, :, None] + position
+
+    def render_viewpoint_data(self, position):
+        data = self.lowres_envmap_renderer.render_viewpoint(np.asarray(position))
+        colormap = data['observations']['color_equirectangular'][..., :3]  # Ignore the alpha channel
+        distancemap = data['observations']['depth_equirectangular']
+        pointmap = self.compute_pointmap(distancemap, position)
+        return ViewpointData(colormap=colormap, distancemap=distancemap, pointmap=pointmap, position=position)
+
+    def extract_cropped_camera(self, projection, color_image, distancemap, pointmap, voxelmap=None):
+        remapper = projections_conversions.RemapProjection(input_projection=self.envmap_projection, output_projection=projection,
+                                                           pixel_jittering_iterations=self.pixel_jittering_iterations, jittering_noise_level=self.jittering_noise_level)
+        cropped_color_image = remapper.convert(
+            color_image, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_WRAP, single_map=False)
+        cropped_distancemap = remapper.convert(
+            distancemap, interpolation=cv2.INTER_NEAREST, borderMode=cv2.BORDER_WRAP, single_map=True)
+        cropped_pointmap = remapper.convert(pointmap, interpolation=cv2.INTER_NEAREST,
+                                            borderMode=cv2.BORDER_WRAP, single_map=True)
+        cropped_voxelmap = (None if voxelmap is None else
+                            remapper.convert(voxelmap, interpolation=cv2.INTER_NEAREST, borderMode=cv2.BORDER_WRAP, single_map=True))
+        # Convert the distance map into a depth map
+        cropped_depthmap = np.asarray(
+            cropped_distancemap / np.linalg.norm(remapper.output_rays, axis=-1), dtype=cropped_distancemap.dtype)
+
+        return cropped_color_image, cropped_depthmap, cropped_pointmap, cropped_voxelmap
+
+def perspective_projection_to_dict(persp_projection, position):
+    """
+    Serialization-like function."""
+    camera_params = dict(camera_intrinsics=projections.colmap_to_opencv_intrinsics(persp_projection.base_projection.K).tolist(),
+                         size=(persp_projection.base_projection.width, persp_projection.base_projection.height),
+                         R_cam2world=persp_projection.R_to_base_projection.T.tolist(),
+                         t_cam2world=position)
+    return camera_params
+
+
+def dict_to_perspective_projection(camera_params):
+    K = projections.opencv_to_colmap_intrinsics(np.asarray(camera_params["camera_intrinsics"]))
+    size = camera_params["size"]
+    R_cam2world = np.asarray(camera_params["R_cam2world"])
+    projection = projections.PerspectiveProjection(K, height=size[1], width=size[0])
+    projection = projections.RotatedProjection(projection, R_to_base_projection=R_cam2world.T)
+    position = camera_params["t_cam2world"]
+    return projection, position

vismatch/third_party/mast3r/dust3r/datasets_preprocess/habitat/habitat_renderer/projections.py

@@ -0,0 +1,151 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Various 3D/2D projection utils, useful to sample virtual cameras.
+# --------------------------------------------------------
+import numpy as np
+
+class EquirectangularProjection:
+    """
+    Convention for the central pixel of the equirectangular map similar to OpenCV perspective model:
+        +X from left to right
+        +Y from top to bottom
+        +Z going outside the camera
+    EXCEPT that the top left corner of the image is assumed to have (0,0) coordinates (OpenCV assumes (-0.5,-0.5))
+    """
+
+    def __init__(self, height, width):
+        self.height = height
+        self.width = width
+        self.u_scaling = (2 * np.pi) / self.width
+        self.v_scaling = np.pi / self.height
+
+    def unproject(self, u, v):
+        """
+        Args:
+            u, v: 2D coordinates
+        Returns:
+            unnormalized 3D rays.
+        """
+        longitude = self.u_scaling * u - np.pi
+        minus_latitude = self.v_scaling * v - np.pi/2
+
+        cos_latitude = np.cos(minus_latitude)
+        x, z = np.sin(longitude) * cos_latitude, np.cos(longitude) * cos_latitude
+        y = np.sin(minus_latitude)
+
+        rays = np.stack([x, y, z], axis=-1)
+        return rays
+
+    def project(self, rays):
+        """
+        Args:
+            rays: Bx3 array of 3D rays.
+        Returns:
+            u, v: tuple of 2D coordinates.
+        """
+        rays = rays / np.linalg.norm(rays, axis=-1, keepdims=True)
+        x, y, z = [rays[..., i] for i in range(3)]
+
+        longitude = np.arctan2(x, z)
+        minus_latitude = np.arcsin(y)
+
+        u = (longitude + np.pi) * (1.0 / self.u_scaling)
+        v = (minus_latitude + np.pi/2) * (1.0 / self.v_scaling)
+        return u, v
+
+
+class PerspectiveProjection:
+    """
+    OpenCV convention:
+    World space:
+        +X from left to right
+        +Y from top to bottom
+        +Z going outside the camera
+    Pixel space:
+        +u from left to right
+        +v from top to bottom
+    EXCEPT that the top left corner of the image is assumed to have (0,0) coordinates (OpenCV assumes (-0.5,-0.5)).
+    """
+
+    def __init__(self, K, height, width):
+        self.height = height
+        self.width = width
+        self.K = K
+        self.Kinv = np.linalg.inv(K)
+
+    def project(self, rays):
+        uv_homogeneous = np.einsum("ik, ...k -> ...i", self.K, rays)
+        uv = uv_homogeneous[..., :2] / uv_homogeneous[..., 2, None]
+        return uv[..., 0], uv[..., 1]
+
+    def unproject(self, u, v):
+        uv_homogeneous = np.stack((u, v, np.ones_like(u)), axis=-1)
+        rays = np.einsum("ik, ...k -> ...i", self.Kinv, uv_homogeneous)
+        return rays
+
+
+class RotatedProjection:
+    def __init__(self, base_projection, R_to_base_projection):
+        self.base_projection = base_projection
+        self.R_to_base_projection = R_to_base_projection
+
+    @property
+    def width(self):
+        return self.base_projection.width
+
+    @property
+    def height(self):
+        return self.base_projection.height
+
+    def project(self, rays):
+        if self.R_to_base_projection is not None:
+            rays = np.einsum("ik, ...k -> ...i", self.R_to_base_projection, rays)
+        return self.base_projection.project(rays)
+
+    def unproject(self, u, v):
+        rays = self.base_projection.unproject(u, v)
+        if self.R_to_base_projection is not None:
+            rays = np.einsum("ik, ...k -> ...i", self.R_to_base_projection.T, rays)
+        return rays
+
+def get_projection_rays(projection, noise_level=0):
+    """
+    Return a 2D map of 3D rays corresponding to the projection.
+    If noise_level > 0, add some jittering noise to these rays.
+    """
+    grid_u, grid_v = np.meshgrid(0.5 + np.arange(projection.width), 0.5 + np.arange(projection.height))
+    if noise_level > 0:
+        grid_u += np.clip(0, noise_level * np.random.uniform(-0.5, 0.5, size=grid_u.shape), projection.width)
+        grid_v += np.clip(0, noise_level * np.random.uniform(-0.5, 0.5, size=grid_v.shape), projection.height)
+    return projection.unproject(grid_u, grid_v)
+
+def compute_camera_intrinsics(height, width, hfov):
+    f = width/2 / np.tan(hfov/2 * np.pi/180)
+    cu, cv = width/2, height/2
+    return f, cu, cv
+
+def colmap_to_opencv_intrinsics(K):
+    """
+    Modify camera intrinsics to follow a different convention.
+    Coordinates of the center of the top-left pixels are by default:
+    - (0.5, 0.5) in Colmap
+    - (0,0) in OpenCV
+    """
+    K = K.copy()
+    K[0, 2] -= 0.5
+    K[1, 2] -= 0.5
+    return K
+
+def opencv_to_colmap_intrinsics(K):
+    """
+    Modify camera intrinsics to follow a different convention.
+    Coordinates of the center of the top-left pixels are by default:
+    - (0.5, 0.5) in Colmap
+    - (0,0) in OpenCV
+    """
+    K = K.copy()
+    K[0, 2] += 0.5
+    K[1, 2] += 0.5
+    return K

vismatch/third_party/mast3r/dust3r/datasets_preprocess/habitat/habitat_renderer/projections_conversions.py

@@ -0,0 +1,45 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Remap data from one projection to an other
+# --------------------------------------------------------
+import numpy as np
+import cv2
+from habitat_renderer import projections
+
+class RemapProjection:
+    def __init__(self, input_projection, output_projection, pixel_jittering_iterations=0, jittering_noise_level=0):
+        """
+        Some naive random jittering can be introduced in the remapping to mitigate aliasing artecfacts.
+        """
+        assert jittering_noise_level >= 0
+        assert pixel_jittering_iterations >= 0
+
+        maps = []
+        # Initial map
+        self.output_rays = projections.get_projection_rays(output_projection)
+        map_u, map_v = input_projection.project(self.output_rays)
+        map_u, map_v = np.asarray(map_u, dtype=np.float32), np.asarray(map_v, dtype=np.float32)
+        maps.append((map_u, map_v))
+
+        for _ in range(pixel_jittering_iterations):
+            # Define multiple mappings using some coordinates jittering to mitigate aliasing effects
+            crop_rays = projections.get_projection_rays(output_projection, jittering_noise_level)
+            map_u, map_v = input_projection.project(crop_rays)
+            map_u, map_v = np.asarray(map_u, dtype=np.float32), np.asarray(map_v, dtype=np.float32)
+            maps.append((map_u, map_v))
+        self.maps = maps
+
+    def convert(self, img, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_WRAP, single_map=False):
+        remapped = []
+        for map_u, map_v in self.maps:
+            res = cv2.remap(img, map_u, map_v, interpolation=interpolation, borderMode=borderMode)
+            remapped.append(res)
+            if single_map:
+                break
+        if len(remapped) == 1:
+            res = remapped[0]
+        else:
+            res = np.asarray(np.mean(remapped, axis=0), dtype=img.dtype)
+        return res

vismatch/third_party/mast3r/dust3r/datasets_preprocess/habitat/preprocess_habitat.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# main executable for preprocessing habitat
+# export METADATA_DIR="/path/to/habitat/5views_v1_512x512_metadata"
+# export SCENES_DIR="/path/to/habitat/data/scene_datasets/"
+# export OUTPUT_DIR="data/habitat_processed"
+# export PYTHONPATH=$(pwd)
+# python preprocess_habitat.py --scenes_dir=$SCENES_DIR --metadata_dir=$METADATA_DIR --output_dir=$OUTPUT_DIR | parallel -j 16
+# --------------------------------------------------------
+import os
+import glob
+import json
+import os
+
+import PIL.Image
+import json
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"  # noqa
+import cv2
+from habitat_renderer import multiview_crop_generator
+from tqdm import tqdm
+
+
+def preprocess_metadata(metadata_filename,
+                        scenes_dir,
+                        output_dir,
+                        crop_resolution=[512, 512],
+                        equirectangular_resolution=None,
+                        fix_existing_dataset=False):
+    # Load data
+    with open(metadata_filename, "r") as f:
+        metadata = json.load(f)
+
+    if metadata["scene_dataset_config_file"] == "":
+        scene = os.path.join(scenes_dir, metadata["scene"])
+        scene_dataset_config_file = ""
+    else:
+        scene = metadata["scene"]
+        scene_dataset_config_file = os.path.join(scenes_dir, metadata["scene_dataset_config_file"])
+    navmesh = None
+
+    # Use 4 times the crop size as resolution for rendering the environment map.
+    max_res = max(crop_resolution)
+
+    if equirectangular_resolution == None:
+        # Use 4 times the crop size as resolution for rendering the environment map.
+        max_res = max(crop_resolution)
+        equirectangular_resolution = (4*max_res, 8*max_res)
+
+    print("equirectangular_resolution:", equirectangular_resolution)
+
+    if os.path.exists(output_dir) and not fix_existing_dataset:
+        raise FileExistsError(output_dir)
+
+    # Lazy initialization
+    highres_dataset = None
+
+    for batch_label, batch in tqdm(metadata["view_batches"].items()):
+        for view_label, view_params in batch.items():
+
+            assert view_params["size"] == crop_resolution
+            label = f"{batch_label}_{view_label}"
+
+            output_camera_params_filename = os.path.join(output_dir, f"{label}_camera_params.json")
+            if fix_existing_dataset and os.path.isfile(output_camera_params_filename):
+                # Skip generation if we are fixing a dataset and the corresponding output file already exists
+                continue
+
+            # Lazy initialization
+            if highres_dataset is None:
+                highres_dataset = multiview_crop_generator.HabitatMultiviewCrops(scene=scene,
+                                                                                 navmesh=navmesh,
+                                                                                 scene_dataset_config_file=scene_dataset_config_file,
+                                                                                 equirectangular_resolution=equirectangular_resolution,
+                                                                                 crop_resolution=crop_resolution,)
+                os.makedirs(output_dir, exist_ok=bool(fix_existing_dataset))
+
+            # Generate a higher resolution crop
+            original_projection, position = multiview_crop_generator.dict_to_perspective_projection(view_params)
+            # Render an envmap at the given position
+            viewpoint_data = highres_dataset.render_viewpoint_data(position)
+
+            projection = original_projection
+            colormap, depthmap, pointmap, _ = highres_dataset.extract_cropped_camera(
+                projection, viewpoint_data.colormap, viewpoint_data.distancemap, viewpoint_data.pointmap)
+
+            camera_params = multiview_crop_generator.perspective_projection_to_dict(projection, position)
+
+            # Color image
+            PIL.Image.fromarray(colormap).save(os.path.join(output_dir, f"{label}.jpeg"))
+            # Depth image
+            cv2.imwrite(os.path.join(output_dir, f"{label}_depth.exr"),
+                        depthmap, [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF])
+            # Camera parameters
+            with open(output_camera_params_filename, "w") as f:
+                json.dump(camera_params, f)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--metadata_dir", required=True)
+    parser.add_argument("--scenes_dir", required=True)
+    parser.add_argument("--output_dir", required=True)
+    parser.add_argument("--metadata_filename", default="")
+
+    args = parser.parse_args()
+
+    if args.metadata_filename == "":
+        # Walk through the metadata dir to generate commandlines
+        for filename in glob.iglob(os.path.join(args.metadata_dir, "**/metadata.json"), recursive=True):
+            output_dir = os.path.join(args.output_dir, os.path.relpath(os.path.dirname(filename), args.metadata_dir))
+            if not os.path.exists(output_dir):
+                commandline = f"python {__file__} --metadata_filename={filename} --metadata_dir={args.metadata_dir} --scenes_dir={args.scenes_dir} --output_dir={output_dir}"
+                print(commandline)
+    else:
+        preprocess_metadata(metadata_filename=args.metadata_filename,
+                            scenes_dir=args.scenes_dir,
+                            output_dir=args.output_dir)

vismatch/third_party/mast3r/dust3r/datasets_preprocess/path_to_root.py

@@ -0,0 +1,13 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# DUSt3R repo root import
+# --------------------------------------------------------
+
+import sys
+import os.path as path
+HERE_PATH = path.normpath(path.dirname(__file__))
+DUST3R_REPO_PATH = path.normpath(path.join(HERE_PATH, '../'))
+# workaround for sibling import
+sys.path.insert(0, DUST3R_REPO_PATH)