Apply Umeyama alignment on a time window

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)

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))

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",