vprdb.core.voxel_grid

 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 nptyping import Float, NDArray, Shape
19
20
21@dataclass
22class VoxelGrid:
23    """Voxel grid with given boundaries and voxel size"""
24
25    min_bounds: NDArray[Shape["3"], Float]
26    max_bounds: NDArray[Shape["3"], Float]
27    voxel_size: float
28
29    def get_voxel_index(
30        self, point: NDArray[Shape["3"], Float]
31    ) -> tuple[int, int, int]:
32        """
33        The method gets the voxel index for a given point
34        Implemented according to the corresponding Open3D method
35        :param point: Point to get its corresponding voxel index
36        :return: Voxel index
37        """
38        ref_coord = (point - self.min_bounds) / self.voxel_size
39        x_index, y_index, z_index = np.floor(ref_coord).astype(int)
40        return x_index, y_index, z_index
41
42    def get_voxel_coordinates(
43        self, point: NDArray[Shape["3"], Float]
44    ) -> tuple[float, float, float]:
45        """
46        The method gets the voxel coordinates for a given point
47        Implemented according to the corresponding Open3D method
48        :param point: Point to get its corresponding voxel coordinates
49        :return: Voxel coordinates
50        """
51        return tuple(
52            self.min_bounds
53            + (np.asarray(self.get_voxel_index(point)) * self.voxel_size)
54        )
55
56    def voxel_down_sample(
57        self, point_cloud: o3d.geometry.PointCloud
58    ) -> o3d.geometry.PointCloud:
59        """
60        Voxel down sampling with bounds given
61        :param point_cloud: Point cloud for down sampling
62        :return: Down sampled point cloud
63        """
64        voxel_down_result, _, _ = point_cloud.voxel_down_sample_and_trace(
65            self.voxel_size, self.min_bounds, self.max_bounds
66        )
67        return voxel_down_result