From cb42cfe725dfc456a8d814bc78160ac770b63769 Mon Sep 17 00:00:00 2001 From: EmmanuelMess Date: Mon, 6 Jan 2025 06:27:49 -0300 Subject: [PATCH] Umeyama on a time window --- evo/core/trajectory.py | 78 +++++++++++++++++++++++++++++++++++++++++ evo/main_traj.py | 9 +++-- evo/main_traj_parser.py | 8 +++++ 3 files changed, 93 insertions(+), 2 deletions(-) diff --git a/evo/core/trajectory.py b/evo/core/trajectory.py index 7367428..ee6757f 100644 --- a/evo/core/trajectory.py +++ b/evo/core/trajectory.py @@ -420,6 +420,84 @@ def speeds(self) -> np.ndarray: for i in range(len(self.positions_xyz) - 1) ]) + def align_on_window(self, traj_ref: 'PoseTrajectory3D', correct_scale: bool = False, + correct_only_scale: bool = False, n: int = -1, + start_time: typing.Optional[float] = None, + end_time: typing.Optional[float] = None) -> geometry.UmeyamaResult: + """ + align to a reference trajectory using Umeyama alignment + :param traj_ref: reference trajectory + :param correct_scale: set to True to adjust also the scale + :param correct_only_scale: set to True to correct the scale, but not the pose + :param n: the number of poses to use, counted from the start (default: all) + :param start_time: the time to start the Umeyama alignment window + (default: start of the trajectory) + :param end_time: the time to end the Umeyama alignment window + (default: end of the trajectory) + :return: the result parameters of the Umeyama algorithm + """ + if start_time is None and end_time is None: + return self.align(traj_ref, correct_scale, correct_only_scale, n) + + if n != -1: + # Cannot have start_time not None or end_time not None, and n != 1 + raise TrajectoryException("start_time or end_time with n is not implemented") + + with_scale = correct_scale or correct_only_scale + if correct_only_scale: + logger.debug("Correcting scale...") + else: + logger.debug(f"Aligning using Umeyama's method... " + f"{'(with scale correction)' if with_scale else ''}") + + relative_timestamps = self.timestamps - np.min(self.timestamps) + + if start_time is None: + start_index = 0 + elif np.all(relative_timestamps < start_time): + logger.warning(f"Align start time ({start_time}s) is after end of trajectory" + f" ({np.max(relative_timestamps)}s), ignoring start time") + start_index = 0 + else: + # Find first value that is less or equal to start_time + start_index = np.flatnonzero(start_time <= relative_timestamps)[0] + logger.debug(f"Start of alignment: in reference {traj_ref.timestamps[start_index]}s, " + f"in trajectory {self.timestamps[start_index]}s") + + if end_time is None: + end_index = self.positions_xyz.shape[0] + elif np.all(relative_timestamps < end_time): + logger.warning(f"Align end time ({end_time}s) is after end of trajectory " + f"({np.max(relative_timestamps)}s), ignoring end time") + end_index = self.timestamps.shape[0] + else: + # Find first value that is greater or equal to end_time + end_index = np.flatnonzero(end_time <= relative_timestamps)[0] + logger.debug(f"End of alignment: in reference {traj_ref.timestamps[end_index]}s, " + f"in trajectory {self.timestamps[end_index]}s") + + if end_index <= start_index: + raise TrajectoryException("alignment is empty") + + r_a, t_a, s = geometry.umeyama_alignment(self.positions_xyz[start_index:end_index, :].T, + traj_ref.positions_xyz[start_index:end_index, :].T, + with_scale) + + if not correct_only_scale: + logger.debug(f"Rotation of alignment:\n{r_a}" + f"\nTranslation of alignment:\n{t_a}") + logger.debug(f"Scale correction: {s}") + + if correct_only_scale: + self.scale(s) + elif correct_scale: + self.scale(s) + self.transform(lie.se3(r_a, t_a)) + else: + self.transform(lie.se3(r_a, t_a)) + + return r_a, t_a, s + def reduce_to_ids( self, ids: typing.Union[typing.Sequence[int], np.ndarray]) -> None: super(PoseTrajectory3D, self).reduce_to_ids(ids) diff --git a/evo/main_traj.py b/evo/main_traj.py index 579ab31..fd772e9 100755 --- a/evo/main_traj.py +++ b/evo/main_traj.py @@ -235,6 +235,10 @@ def run(args): if args.n_to_align != -1 and not (args.align or args.correct_scale): die("--n_to_align is useless without --align or/and --correct_scale") + if args.n_to_align != -1 and (args.start_t_to_align is not None or args.end_t_to_align is not None): + die("--start_t_to_align or --end_t_to_align with --n_to_align is not implemented") + if (args.start_t_to_align is not None or args.end_t_to_align is not None) and not (args.align or args.correct_scale): + die("--start_t_to_align and --end_t_to_align are useless without --align or/and --correct_scale") # TODO: this is fugly, but is a quick solution for remembering each synced # reference when plotting pose correspondences later... @@ -257,10 +261,11 @@ def run(args): if args.align or args.correct_scale: logger.debug(SEP) logger.debug("Aligning {} to reference.".format(name)) - trajectories[name].align( + trajectories[name].align_on_window( ref_traj_tmp, correct_scale=args.correct_scale, correct_only_scale=args.correct_scale and not args.align, - n=args.n_to_align) + n=args.n_to_align, start_time=args.start_t_to_align, + end_time=args.end_t_to_align) if args.align_origin: logger.debug(SEP) logger.debug("Aligning {}'s origin to reference.".format(name)) diff --git a/evo/main_traj_parser.py b/evo/main_traj_parser.py index 94beb08..5b5771f 100644 --- a/evo/main_traj_parser.py +++ b/evo/main_traj_parser.py @@ -20,6 +20,14 @@ def parser() -> argparse.ArgumentParser: "--n_to_align", help="the number of poses to use for Umeyama alignment, " "counted from the start (default: all)", default=-1, type=int) + algo_opts.add_argument( + "--start_t_to_align", + help="the start of the time window to use for Umeyama alignment, " + "in seconds relative to the first timestamp of the file", default=None, type=float) + algo_opts.add_argument( + "--end_t_to_align", + help="the end of the time window to use for Umeyama alignment, " + "in seconds relative to the first timestamp of the file", default=None, type=float) algo_opts.add_argument( "--sync", help="associate trajectories via matching timestamps - requires --ref",