From bc01d0d55140dc3a9342a7abd5e725edd02309f5 Mon Sep 17 00:00:00 2001 From: Mitchell Date: Thu, 30 Jan 2025 11:29:52 -0500 Subject: [PATCH] Add converters to evo to perform trajectory alignment --- navlie/utils/alignment.py | 163 +++++++++++++++++++++++++++++++++++ setup.py | 1 + tests/unit/test_alignment.py | 86 ++++++++++++++++++ 3 files changed, 250 insertions(+) create mode 100644 navlie/utils/alignment.py create mode 100644 tests/unit/test_alignment.py diff --git a/navlie/utils/alignment.py b/navlie/utils/alignment.py new file mode 100644 index 00000000..03ec5c40 --- /dev/null +++ b/navlie/utils/alignment.py @@ -0,0 +1,163 @@ +""" +Utilities for interfacing with the evo package for trajectory alignment and +evaluation. +""" + +import numpy as np +import typing +from pymlg import SE3, SO3 +from navlie.lib.states import SE3State + +from evo.core import trajectory +from evo.core import sync + + +def state_list_to_evo_traj( + state_list: typing.List[SE3State], +) -> trajectory.PoseTrajectory3D: + """Converts a list of SE3States to an evo trajectory. + + Parameters + ---------- + state_list : typing.List[SE3State] + The list of SE3States to convert + + Returns + ------- + trajectory.PoseTrajectory3D + The converted evo trajectory + """ + positions = np.array([state.position for state in state_list]) + quats = np.array( + [SO3.to_quat(x.attitude, order="wxyz").ravel() for x in state_list] + ) + stamps = np.array([x.stamp for x in state_list]) + traj = trajectory.PoseTrajectory3D(positions, quats, stamps) + return traj + + +def evo_traj_to_state_list( + traj: trajectory.PoseTrajectory3D, +) -> typing.List[SE3State]: + """Converts an evo trajectory to a list of SE3State. + + Parameters + ---------- + traj : trajectory.PoseTrajectory3D + The evo pose trajectory to convert + + Returns + ------- + typing.List[SE3State] + The converted list of SE3State + """ + state_list = [] + for i in range(len(traj.timestamps)): + pose_mat = SE3.from_components( + SO3.from_quat(traj.orientations_quat_wxyz[i, :], order="wxyz"), + traj._positions_xyz[i, :], + ) + state = SE3State( + value=pose_mat, + stamp=traj.timestamps[i], + ) + state.stamp = traj.timestamps[i] + state_list.append(state) + return state_list + + +def associate_and_align_trajectories( + traj_ref_list: typing.List[SE3State], + traj_est_list: typing.List[SE3State], + max_diff: float = 0.02, + offset: float = 0.0, + align: bool = True, + correct_scale: bool = False, + n_to_align: int = -1, + verbose: bool = False, +) -> typing.Tuple[ + typing.List[SE3State], typing.List[SE3State], typing.Dict[str, np.ndarray] +]: + """Associates the stamps of two trajectories and aligns the + estimated trajectory to the reference trajectory. + + Parameters + ---------- + traj_ref : typing.List[SE3State] + Reference trajectory, a list of SE3State + traj_est : typing.List[SE3State] + Estimated trajectory to be aligned to the reference, a list of SE3State + max_diff : float, optional + The maximum allowable difference in timestamps between the estimate + and the reference trajectory, by default 0.02 + offset : float, optional + optional time offset of the second trajectory, by default 0.0 + align: bool, optional + Whether to align the trajectories, by default True + correct_scale: bool, optional + Whether to correct the scale, by default False. If correct_scale is + false, the alignment transformation is an element of SE(3), otherwise, + it is an element of Sim(3). + n_to_align: int, optional + The number of poses to use for alignment, by default -1. If n_to_align + is -1, all poses are used. + verbose : bool, optional + Verbosity flag, by default False + + Returns + ------- + typing.List[SE3State] + The reference trajectory, with timestamps at the same times as + the estimated trajectory. + typing.List[SE3State] + The aligned estimated trajectory. + typing.Dict[str, np.ndarray] + A dictionary containing the transformation information. The keys are: + - "rotation": The DCM of the alignment transformation + - "position": The transformation component of the alignment transformation + - "scale": The scale component alignment transformation. + """ + # Convert trajectories to evo, and associate the timestamps + traj_ref = state_list_to_evo_traj(traj_ref_list) + traj_est = state_list_to_evo_traj(traj_est_list) + traj_ref_sync, traj_est_sync = sync.associate_trajectories( + traj_ref, + traj_est, + max_diff=max_diff, + offset_2=offset, + ) + + if verbose: + print(f"Number of matching states: {len(traj_ref_sync.timestamps)}") + print(f"Length of reference trajectory: {len(traj_ref_list)}") + print(f"Length of estimated trajectory: {len(traj_est_list)}") + + # If we are not aligning, return the synced trajectories + if not align: + traj_ref_sync = evo_traj_to_state_list(traj_ref_sync) + traj_est_sync = evo_traj_to_state_list(traj_est_sync) + transformation_dict = { + "rotation": np.identity(3), + "position": np.zeros(3), + "scale": 1.0, + } + return traj_ref_sync, traj_est_sync, transformation_dict + + # Otherwise, align the trajectories + traj_est_sync: trajectory.PoseTrajectory3D + R, t, s = traj_est_sync.align( + traj_ref_sync, + correct_scale=correct_scale, + n=n_to_align, + ) + + # Conver back to navlie types + traj_est_aligned = evo_traj_to_state_list(traj_est_sync) + traj_ref_aligned = evo_traj_to_state_list(traj_ref_sync) + transformation_dict = { + "rotation": R, + "position": t, + "scale": s, + } + + return traj_ref_aligned, traj_est_aligned, transformation_dict diff --git a/setup.py b/setup.py index 7b71551b..2cc0b778 100644 --- a/setup.py +++ b/setup.py @@ -18,6 +18,7 @@ "pymlg @ git+https://github.com/decargroup/pymlg@main", "tqdm>=4.64.1", "seaborn>=0.11.2", + "evo>=1.31.0" ], url="https://github.com/decargroup/navlie", classifiers=[ diff --git a/tests/unit/test_alignment.py b/tests/unit/test_alignment.py new file mode 100644 index 00000000..3309b852 --- /dev/null +++ b/tests/unit/test_alignment.py @@ -0,0 +1,86 @@ +import typing +import numpy as np +from navlie.lib.datasets import SimulatedPoseRangingDataset +from navlie.lib.states import SE3State +from navlie.utils import plot_poses + +from pymlg import SO3, SE3 +import matplotlib.pyplot as plt + +from navlie.utils.alignment import associate_and_align_trajectories, state_list_to_evo_traj, evo_traj_to_state_list + +def test_conversion_to_evo(): + """Converts a list of SE3State to an evo trajectory and back.""" + dataset = SimulatedPoseRangingDataset() + poses_1 = dataset.get_ground_truth() + evo_traj = state_list_to_evo_traj(poses_1) + poses_2 = evo_traj_to_state_list(evo_traj) + + for pose_1, pose_2 in zip(poses_1, poses_2): + assert np.allclose(pose_1.position, pose_2.position, atol=1e-5) + assert np.allclose(pose_1.attitude, pose_2.attitude, atol=1e-5) + assert np.allclose(pose_1.stamp, pose_2.stamp, atol=1e-5) + +def test_associate_and_align_trajectories(): + """Test utilizing evo to associate and align two trajectories.""" + # Load in the sample EuRoC dataset + # Create an initial trajectory + dataset = SimulatedPoseRangingDataset() + poses_1 = dataset.get_ground_truth() + + # Create a second trajectory that is a similarity transformation of the first + scale = 0.9 + rot = SO3.Exp(np.array([0.01, 0.01, 1.2])) + pos = np.array([1.0, 2.0, 0.5]) + + poses_2: typing.List[SE3State] = [] + for pose in poses_1: + new_position = scale * rot @ pose.position + pos + new_att = rot @ pose.attitude + new_pose = SE3State( + value=SE3.from_components(new_att, new_position), + stamp=pose.stamp, + direction=pose.direction, + ) + poses_2.append(new_pose) + + fig, ax = plot_poses( + poses_1, line_color="tab:blue", step=None, label="Original Trajectory" + ) + fig, ax = plot_poses( + poses_2, ax=ax, line_color="tab:red", step=None, label="Transformed Trajectory" + ) + ax.set_xlabel("x (m)") + ax.set_ylabel("y (m)") + ax.set_zlabel("z (m)") + + # Try aligning the first trajectory to the second + traj_ref, aligned_traj_sim3, transformation_info = associate_and_align_trajectories( + poses_2, + poses_1, + correct_scale=True, + align=True, + verbose=True + ) + + # Check that the computed transformation matches the true one + assert np.allclose(transformation_info["scale"], scale, atol=1e-5) + assert np.allclose(transformation_info["rotation"], rot, atol=1e-5) + assert np.allclose(transformation_info["position"], pos, atol=1e-5) + + # Verify visually + fig, ax = plot_poses( + aligned_traj_sim3, + ax=ax, + line_color="tab:green", + step=None, + label="Aligned Trajectory (Sim3)", + ) + + ax.legend() + fig.tight_layout() + plt.show() + +if __name__ == "__main__": + test_conversion_to_evo() + test_associate_and_align_trajectories() \ No newline at end of file