vprdb.core.utils

  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 nptyping import Float, NDArray, Shape
 18
 19from vprdb.core.database import Database
 20from vprdb.core.voxel_grid import VoxelGrid
 21
 22
 23def calculate_point_cloud_coverage(
 24    query_pcd: o3d.geometry.PointCloud,
 25    db_pcd: o3d.geometry.PointCloud,
 26    voxel_grid: VoxelGrid,
 27) -> float:
 28    """
 29    Calculates coverage for query point cloud by database point cloud
 30    :param query_pcd: query point cloud
 31    :param db_pcd: database point cloud
 32    :param voxel_grid: voxel grid for down sampling
 33    :return: coverage ∈ [0; 1]
 34    """
 35    query_pcd = voxel_grid.voxel_down_sample(query_pcd)
 36    db_pcd = voxel_grid.voxel_down_sample(db_pcd)
 37    united_map = voxel_grid.voxel_down_sample(query_pcd + db_pcd)
 38
 39    query_size = len(query_pcd.points)
 40    intersection_size = query_size + len(db_pcd.points) - len(united_map.points)
 41    return intersection_size / query_size
 42
 43
 44def calculate_iou(
 45    pcd_1: o3d.geometry.PointCloud,
 46    pcd_2: o3d.geometry.PointCloud,
 47    voxel_grid: VoxelGrid,
 48) -> float:
 49    """
 50    Calculates IoU for two point clouds
 51    :param pcd_1: First point cloud
 52    :param pcd_2: Second point cloud
 53    :param voxel_grid: Voxel grid for down sampling
 54    :return: IoU ∈ [0; 1]
 55    """
 56    pcd_1 = voxel_grid.voxel_down_sample(pcd_1)
 57    pcd_2 = voxel_grid.voxel_down_sample(pcd_2)
 58    united_map = voxel_grid.voxel_down_sample(pcd_1 + pcd_2)
 59
 60    united_map_size = len(united_map.points)
 61    intersection_size = len(pcd_1.points) + len(pcd_2.points) - united_map_size
 62    iou = intersection_size / united_map_size
 63    return iou
 64
 65
 66def match_two_databases(
 67    source_db: Database, target_db: Database, voxel_grid: VoxelGrid
 68) -> list[int]:
 69    """
 70    Builds matches between two databases. Databases should have the same coordinate system
 71    :param source_db: The database for which matches are being built
 72    :param target_db: The database on the base of which the matches will be built
 73    :param voxel_grid: VoxelGrid for down sampling
 74    :return: Output matches.
 75    The i-th value from the list is the index of the frame from the target database,
 76    which is matched with the i-th source frame
 77    """
 78    matches = []
 79    for pose, pcd_raw in zip(source_db.trajectory, source_db.point_clouds):
 80        cur_coverages = []
 81        for target_pose, target_pcd in zip(
 82            target_db.trajectory, target_db.point_clouds
 83        ):
 84            pcd_query = pcd_raw.point_cloud.transform(pose)
 85            pcd_db = target_pcd.point_cloud.transform(target_pose)
 86            coverage = calculate_point_cloud_coverage(pcd_query, pcd_db, voxel_grid)
 87            cur_coverages.append(coverage)
 88        best_match = np.argmax(cur_coverages)
 89        matches.append(int(best_match))
 90    return matches
 91
 92
 93def find_bounds_for_multiple_databases(
 94    databases: list[Database],
 95) -> tuple[NDArray[Shape["3"], Float], NDArray[Shape["3"], Float]]:
 96    """
 97    Finds bounds for multiple databases
 98    :param databases: List of databases for finding bounds
 99    :return: Min and max bounds
100    """
101    min_bounds_global, max_bounds_global = [], []
102    for db in databases:
103        min_bounds, max_bounds = db.bounds
104        min_bounds_global.append(min_bounds)
105        max_bounds_global.append(max_bounds)
106    min_bounds_global, max_bounds_global = np.asarray(min_bounds_global), np.asarray(
107        max_bounds_global
108    )
109    min_bound = np.amin(min_bounds_global, axis=0)
110    max_bound = np.amax(max_bounds_global, axis=0)
111    return min_bound, max_bound