File size: 3,288 Bytes
744b072 d7e30ae 744b072 d7e30ae 744b072 d7e30ae 744b072 d7e30ae 744b072 d7e30ae 744b072 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import numpy as np
from scipy.spatial.transform import Rotation as R
import sys
import os
import pandas as pd
def get_colmap_keyframes(images_file, number_of_header_lines, verbose=False):
print(f"get_colmap_keyframes: {images_file}")
image_id = []
q_wc_xyzw = []
t_wc = []
with open(f"{images_file}", 'r') as file:
# Skip the header lines
for _ in range(number_of_header_lines):
file.readline()
while True:
line1 = file.readline()
if not line1:
break
elements = line1.split()
IMAGE_ID = int(elements[0])
image_id.append(IMAGE_ID)
QW = float(elements[1])
QX = float(elements[2])
QY = float(elements[3])
QZ = float(elements[4])
TX = float(elements[5])
TY = float(elements[6])
TZ = float(elements[7])
t_cw_i = np.array([TX, TY, TZ])
q_wc_i = R.from_quat([QX, QY, QZ, QW]).inv()
R_wc_i = q_wc_i.as_matrix()
q_wc_xyzw.append([q_wc_i.as_quat()[0], q_wc_i.as_quat()[1], q_wc_i.as_quat()[2], q_wc_i.as_quat()[3]])
t_wc.append(-R_wc_i @ t_cw_i)
file.readline()
image_id = np.array(image_id)
q_wc_xyzw = np.array(q_wc_xyzw)
t_wc = np.array(t_wc)
sorted_indices = image_id.argsort()
image_id = image_id[sorted_indices]
q_wc_xyzw = q_wc_xyzw[sorted_indices]
t_wc = t_wc[sorted_indices]
q_wc_xyzw_corrected = q_wc_xyzw.copy()
for i in range(1, len(q_wc_xyzw_corrected)):
dot_product = np.dot(q_wc_xyzw_corrected[i - 1], q_wc_xyzw_corrected[i])
if dot_product < 0:
q_wc_xyzw_corrected[i] = -q_wc_xyzw_corrected[i]
return image_id, t_wc, q_wc_xyzw_corrected
def write_trajectory_tum_format(file_name, image_ts, t_wc, q_wc_xyzw):
print(f"writeTrajectoryTUMformat: {file_name}")
data = np.hstack((image_ts.reshape(-1, 1), t_wc, q_wc_xyzw))
data = data[data[:, 0].argsort()]
with open(file_name, 'w', newline='') as file:
file.write('timestamp,tx,ty,tz,qx,qy,qz,qw\n')
for row in data:
file.write(','.join(f'{x:.15f}' for x in row) + '\n')
def get_timestamps(files_path, rgb_file):
print(f"getTimestamps: {os.path.join(files_path, rgb_file)}")
df = pd.read_csv(rgb_file)
ts = df['ts_rgb0 (s)'].to_list()
return ts
if __name__ == "__main__":
sequence_path = sys.argv[1]
exp_folder = sys.argv[2]
exp_id = sys.argv[3]
verbose = bool(int(sys.argv[4]))
rgb_file = sys.argv[5]
images_file = os.path.join(exp_folder, f'colmap_{exp_id}', 'images.txt')
number_of_header_lines = 4
image_id, t_wc, q_wc_xyzw = get_colmap_keyframes(images_file, number_of_header_lines, verbose)
image_ts = np.array(get_timestamps(sequence_path, rgb_file))
timestamps = []
for id in image_id:
timestamps.append(float(image_ts[id-1]))
timestamps = np.array(timestamps)
keyFrameTrajectory_txt = os.path.join(exp_folder, exp_id + '_KeyFrameTrajectory' + '.csv')
write_trajectory_tum_format(keyFrameTrajectory_txt, timestamps, t_wc, q_wc_xyzw)
|