vprdb.io.read_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 mrob
15import numpy as np
16
17from nptyping import Float, NDArray, Shape
18from pathlib import Path
19
20from vprdb.core import Database
21
22
23def __load_poses(poses_path: Path, with_timestamps: bool):
24    poses_quat = []
25    with open(poses_path, "r") as file:
26        for line in file:
27            poses_quat.append([float(i) for i in line.strip().split(" ")])
28
29    poses = []
30    for pose in poses_quat:
31        t = pose[1:4] if with_timestamps else pose[:3]
32        R = pose[4:8] if with_timestamps else pose[3:7]
33        R = mrob.geometry.quat_to_so3(R)
34        T = np.eye(4)
35        T[:3, :3] = R
36        T[:3, 3] = t
37        poses.append(T)
38
39    return poses
40
41
42def __read_dir(dir_path: Path) -> list[Path]:
43    return sorted(list(dir_path.iterdir()))
44
45
46def read_dataset_from_depth(
47    path_to_dataset: Path,
48    color_dir: str,
49    depth_dir: str,
50    trajectory_file_name: str,
51    intrinsics: NDArray[Shape["3, 3"], Float],
52    depth_scale: int,
53    with_timestamps: bool = True,
54) -> Database:
55    """
56    Reads dataset from given directory
57    :param path_to_dataset: The name of the directory to read
58    :param color_dir: The name of the directory with color images
59    :param depth_dir: The name of the directory with depth images
60    :param trajectory_file_name: The name of file with trajectory
61    :param intrinsics: NumPy array with camera intrinsics
62    :param depth_scale: Depth scale
63    :param with_timestamps: Indicates that a trajectory with timestamps is given
64    """
65    rgb_images = __read_dir(path_to_dataset / color_dir)
66    depth_images = __read_dir(path_to_dataset / depth_dir)
67
68    traj = __load_poses(path_to_dataset / trajectory_file_name, with_timestamps)
69
70    database = Database.from_depth_images(
71        rgb_images, depth_images, depth_scale, intrinsics, traj
72    )
73    return database
74
75
76def read_dataset_from_point_clouds(
77    path_to_dataset: Path,
78    color_dir: str,
79    point_clouds_dir: str,
80    trajectory_file_name: str,
81    with_timestamps: bool = True,
82) -> Database:
83    """
84    Reads dataset from given directory
85    :param path_to_dataset: The name of the directory to read
86    :param color_dir: The name of the directory with color images
87    :param point_clouds_dir: The name of the directory with point clouds
88    :param trajectory_file_name: The name of file with trajectory
89    :param with_timestamps: Indicates that a trajectory with timestamps is given
90    """
91    rgb_images = __read_dir(path_to_dataset / color_dir)
92    point_clouds = __read_dir(path_to_dataset / point_clouds_dir)
93
94    traj = __load_poses(path_to_dataset / trajectory_file_name, with_timestamps)
95
96    database = Database.from_point_clouds(rgb_images, point_clouds, traj)
97    return database