vprdb.core.database

  1#  Copyright (c) 2023, Ivan Moskalenko, Anastasiia Kornilova
  2#
  3#  Licensed under the Apache License, Version 2.0 (the "License");
  4#  you may not use this file except in compliance with the License.
  5#  You may obtain a copy of the License at
  6#
  7#      http://www.apache.org/licenses/LICENSE-2.0
  8#
  9#  Unless required by applicable law or agreed to in writing, software
 10#  distributed under the License is distributed on an "AS IS" BASIS,
 11#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12#  See the License for the specific language governing permissions and
 13#  limitations under the License.
 14import numpy as np
 15import open3d as o3d
 16
 17from dataclasses import dataclass
 18from functools import cached_property
 19from nptyping import Float, NDArray, Shape
 20from pathlib import Path
 21
 22from vprdb.core.voxel_grid import VoxelGrid
 23from vprdb.providers import ColorImageProvider, DepthImageProvider, PointCloudProvider
 24
 25
 26@dataclass(frozen=True)
 27class Database:
 28    """
 29    Class for easy storage and processing of color images, depth images
 30    and trajectory and their further use for the VPR task
 31    """
 32
 33    color_images: list[ColorImageProvider]
 34    point_clouds: list[DepthImageProvider | PointCloudProvider]
 35    trajectory: list[NDArray[Shape["4, 4"], Float]]
 36
 37    def __post_init__(self):
 38        if not (len(self.trajectory) == len(self.point_clouds) == len(self.trajectory)):
 39            raise ValueError(
 40                "Trajectory, RGB images and point clouds should have equal length"
 41            )
 42
 43    @classmethod
 44    def from_depth_images(
 45        cls,
 46        color_images_paths: list[Path],
 47        depth_images_paths: list[Path],
 48        depth_scale: int,
 49        intrinsics: NDArray[Shape["3, 3"], Float],
 50        trajectory: list[NDArray[Shape["4, 4"], Float]],
 51    ):
 52        """
 53        The method allows to construct the Database from scanning sequence
 54        :param color_images_paths: List of paths to color images
 55        :param depth_images_paths: List of paths to depth images
 56        :param depth_scale: Depth scale for transforming depth image into point clouds
 57        :param intrinsics: Intrinsic camera parameters
 58        :param trajectory: List of camera poses
 59        :return: Constructed database
 60        """
 61        color_images_providers = [
 62            ColorImageProvider(path_to_image) for path_to_image in color_images_paths
 63        ]
 64        depth_images_providers = [
 65            DepthImageProvider(path_to_image, intrinsics, depth_scale)
 66            for path_to_image in depth_images_paths
 67        ]
 68        return cls(color_images_providers, depth_images_providers, trajectory)
 69
 70    @classmethod
 71    def from_point_clouds(
 72        cls,
 73        color_images_paths: list[Path],
 74        point_clouds_paths: list[Path],
 75        trajectory: list[NDArray[Shape["4, 4"], Float]],
 76    ):
 77        """
 78        The method allows to construct the Database from point clouds, color images and trajectory
 79        :param color_images_paths: List of paths to color images
 80        :param point_clouds_paths: List of paths to point clouds
 81        :param trajectory: List of camera poses
 82        :return: Constructed database
 83        """
 84        color_images_providers = [
 85            ColorImageProvider(path_to_image) for path_to_image in color_images_paths
 86        ]
 87        point_clouds_providers = [
 88            PointCloudProvider(path_to_pcd) for path_to_pcd in point_clouds_paths
 89        ]
 90        return cls(color_images_providers, point_clouds_providers, trajectory)
 91
 92    def __len__(self):
 93        return len(self.trajectory)
 94
 95    @cached_property
 96    def bounds(
 97        self,
 98    ) -> tuple[NDArray[Shape["3"], Float], NDArray[Shape["3"], Float]]:
 99        """
100        Gets bounds of the DB scene
101        :return: Min and max bounds of the scene
102        """
103        min_bounds = []
104        max_bounds = []
105
106        for i in range(len(self)):
107            pcd = self.point_clouds[i].point_cloud.transform(self.trajectory[i])
108            min_bounds.append(pcd.get_min_bound())
109            max_bounds.append(pcd.get_max_bound())
110
111        min_bound = np.amin(np.asarray(min_bounds), axis=0)
112        max_bound = np.amax(np.asarray(max_bounds), axis=0)
113        return min_bound, max_bound
114
115    def build_sparse_map(
116        self,
117        voxel_grid: VoxelGrid,
118        down_sample_step: int,
119    ) -> o3d.geometry.PointCloud:
120        """
121        Builds sparse map of the whole DB scene
122        :param voxel_grid: Voxel grid for down sampling
123        :param down_sample_step: Voxel down sampling step for reducing RAM usage
124        :return: Resulting point cloud of the scene
125        """
126        map_pcd = o3d.geometry.PointCloud()
127        for i, (pose, pcd_raw) in enumerate(zip(self.trajectory, self.point_clouds)):
128            pcd = pcd_raw.point_cloud.transform(pose)
129            map_pcd += pcd
130            if i % down_sample_step == 0:
131                map_pcd = voxel_grid.voxel_down_sample(map_pcd)
132        return voxel_grid.voxel_down_sample(map_pcd)