Spaces:
Sleeping
Sleeping
File size: 2,885 Bytes
20ae9ff |
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 |
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
from evo.core import sync
from evo.core.trajectory import PoseTrajectory3D
from evo.tools import plot
from plyfile import PlyData, PlyElement
def plot_trajectory(pred_traj, gt_traj=None, title="", filename="", align=True, correct_scale=True):
assert isinstance(pred_traj, PoseTrajectory3D)
if gt_traj is not None:
assert isinstance(gt_traj, PoseTrajectory3D)
gt_traj, pred_traj = sync.associate_trajectories(gt_traj, pred_traj)
if align:
pred_traj.align(gt_traj, correct_scale=correct_scale)
plot_collection = plot.PlotCollection("PlotCol")
fig = plt.figure(figsize=(8, 8))
plot_mode = plot.PlotMode.xz # ideal for planar movement
ax = plot.prepare_axis(fig, plot_mode)
ax.set_title(title)
if gt_traj is not None:
plot.traj(ax, plot_mode, gt_traj, '--', 'gray', "Ground Truth")
plot.traj(ax, plot_mode, pred_traj, '-', 'blue', "Predicted")
plot_collection.add_figure("traj (error)", fig)
plot_collection.export(filename, confirm_overwrite=False)
plt.close(fig=fig)
print(f"Saved {filename}")
def save_output_for_COLMAP(name: str, traj: PoseTrajectory3D, points: np.ndarray, colors: np.ndarray, fx, fy, cx, cy, H=480, W=640):
""" Saves the sparse point cloud and camera poses such that it can be opened in COLMAP """
colmap_dir = Path(name)
colmap_dir.mkdir(exist_ok=True)
scale = 10 # for visualization
# images
images = ""
traj = PoseTrajectory3D(poses_se3=list(map(np.linalg.inv, traj.poses_se3)), timestamps=traj.timestamps)
for idx, (x,y,z), (qw, qx, qy, qz) in zip(range(1,traj.num_poses+1), traj.positions_xyz*scale, traj.orientations_quat_wxyz):
images += f"{idx} {qw} {qx} {qy} {qz} {x} {y} {z} 1\n\n"
(colmap_dir / "images.txt").write_text(images)
# points
points3D = ""
colors_uint = (colors * 255).astype(np.uint8).tolist()
for i, (p,c) in enumerate(zip((points*scale).tolist(), colors_uint), start=1):
points3D += f"{i} " + ' '.join(map(str, p + c)) + " 0.0 0 0 0 0 0 0\n"
(colmap_dir / "points3D.txt").write_text(points3D)
# camera
(colmap_dir / "cameras.txt").write_text(f"1 PINHOLE {W} {H} {fx} {fy} {cx} {cy}")
print(f"Saved COLMAP-compatible reconstruction in {colmap_dir.resolve()}")
def save_ply(name: str, points: np.ndarray, colors: np.ndarray):
points_ply = np.array([(x,y,z,r,g,b) for (x,y,z),(r,g,b) in zip(points, colors)],
dtype=[('x', '<f4'), ('y', '<f4'), ('z', '<f4'),('red', 'u1'), ('green', 'u1'),('blue', 'u1')])
el = PlyElement.describe(points_ply, 'vertex',{'some_property': 'f8'},{'some_property': 'u4'})
PlyData([el], text=True).write(f"{name}.ply")
print(f"Saved {name}.ply")
|