vprdb.io.export_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
16import shutil
17
18from pathlib import Path
19
20from vprdb.core import Database
21
22
23def __poses_to_txt(poses, filename_to_write):
24    with open(filename_to_write, "w") as trajectory_file:
25        for pose in poses:
26            R = pose[:3, :3]
27            t = pose[:3, 3]
28            quat = mrob.geometry.so3_to_quat(R)
29            pose_string = " ".join(np.concatenate((t, quat)).astype(str))
30            trajectory_file.write(f"{pose_string}\n")
31
32
33def export(
34    database: Database,
35    path_to_export: Path,
36    color_dir: str,
37    point_clouds_dir: str,
38    trajectory_file_name: str,
39):
40    """
41    Exports Database to hard drive
42    :param database: Database for exporting
43    :param path_to_export: Directory for exporting. Will be created if it does not exist
44    :param color_dir: Directory name for saving color images
45    :param point_clouds_dir: Directory name for saving depth images / PCDs
46    :param trajectory_file_name: File name for saving the trajectory
47    """
48    path_to_color = path_to_export / color_dir
49    path_to_point_clouds = path_to_export / point_clouds_dir
50    path_to_color.mkdir(parents=True, exist_ok=False)
51    path_to_point_clouds.mkdir(exist_ok=False)
52
53    for rgb_image in database.color_images:
54        shutil.copyfile(rgb_image.path, path_to_color / rgb_image.path.name)
55
56    for point_cloud in database.point_clouds:
57        shutil.copyfile(point_cloud.path, path_to_point_clouds / point_cloud.path.name)
58
59    __poses_to_txt(database.trajectory, path_to_export / trajectory_file_name)