Ball velocity filtering and small bug fix

bin/pam_ball_trajectories.py

@@ -66,7 +66,7 @@ def _add_tennicam(hdf5_path: pathlib.Path, group_name: str):
     with bt.MutableRecordedBallTrajectories(path=hdf5_path) as rbt:
         if group_name in rbt.get_groups():
             raise ValueError("group {} already present in the file")
-    nb_added = rbt.add_tennicam_trajectories(group_name, pathlib.Path.cwd())
+        nb_added = rbt.add_tennicam_trajectories(group_name, pathlib.Path.cwd())
     logging.info("added {} trajectories".format(nb_added))
 
 

python/context/ball_trajectories.py

@@ -8,6 +8,8 @@
 from __future__ import annotations
 import typing
 import nptyping as npt
+from typing import Union, Tuple
+
 
 import random
 import math
@@ -17,7 +19,8 @@
 
 import o80
 import pam_configuration
-import tennicam_client
+
+from context import LowPassFilter
 
 
 assert int(npt.__version__[0]) >= 2, "Need nptyping >=2."
@@ -83,7 +86,7 @@
     return stamps, positions
 
 
-def to_duration_trajectory(input: StampedTrajectory) -> DurationTrajectory:
+def to_duration_trajectory(input: StampedTrajectory, vel_filter_window_size: int = 1) -> DurationTrajectory:
     """
     Converts a StampedTrajectories to a DurationTrajectory.
     The velocities are estimated by performing finite differences.
@@ -92,6 +95,12 @@
     dp = np.diff(input[1], axis=0)
     velocities = (dp.T / (dt * 1e-6)).T
     positions = input[1][:-1, :]
+
+    #velocity filtering
+    if vel_filter_window_size > 1:
+        filter_lp = [LowPassFilter(vel_filter_window_size), LowPassFilter(vel_filter_window_size), LowPassFilter(vel_filter_window_size)]
+        velocities = [[filter_lp[i].get(v[i]) for i in range(3)] for v in velocities]
+
     return dt, positions, velocities
 
 
@@ -130,7 +139,7 @@
     _TIME_STAMPS = "time_stamps"
     _TRAJECTORY = "trajectory"
 
     def __init__(self, path: pathlib.Path = None, file_mode: str = "r"):
         if path is None:
             path = self.get_default_path()
         self._f = h5py.File(path, file_mode)
@@ -247,7 +256,7 @@
     "r+" mode.
     """
 
     def __init__(self, path: pathlib.Path = None):
         super().__init__(path, file_mode="r+")
 
     def rm_group(self, group: str) -> None:
@@ -290,6 +299,8 @@
         The number of trajectories added to the file.
         """
 
+        import tennicam_client
+
         def _read_trajectory(tennicam_file: pathlib.Path) -> StampedTrajectory:
             """
             Parse the file and returned the corresponding
@@ -445,7 +456,7 @@
       /opt/mpi-is/pam/context/ball_trajectories.hdf5
     """
 
     def __init__(self, group: str, hdf5_path: pathlib.Path = None):
         if hdf5_path is None:
             hdf5_path = RecordedBallTrajectories.get_default_path()
 
@@ -455,7 +466,7 @@
             self._data: typing.Dict[
                 int, StampedTrajectory
             ] = rbt.get_stamped_trajectories(group, direct=True)
-
+            
     def size(self) -> int:
         """
         Returns the number of trajectories that have been loaded.
@@ -475,12 +486,15 @@
         """
         return self._data[index]
 
-    def random_trajectory(self) -> StampedTrajectory:
+    def random_trajectory(self, return_index: bool = False) -> Union[StampedTrajectory, Tuple[StampedTrajectory, int]]:
         """
         Returns one of the trajectory, randomly selected.
+        If return_index is True, also returns the selected index.
         """
         index = random.choice(list(range(len(self._data.keys()))))
-        return self._data[index]
+        trajectory = self._data[index]
+
+        return (trajectory, index) if return_index else trajectory
 
     def get_different_random_trajectories(
         self, nb_trajectories: int
@@ -499,22 +513,22 @@
         return [self._data[index] for index in indexes[:nb_trajectories]]
 
     @staticmethod
-    def to_duration(input: StampedTrajectory) -> DurationTrajectory:
+    def to_duration(input: StampedTrajectory, vel_filter_window_size: int  = 1) -> DurationTrajectory:
         """
         Returns a corresponding duration trajectory
         """
-        return to_duration_trajectory(input)
+        return to_duration_trajectory(input, vel_filter_window_size)
 
     @classmethod
     def iterate(
-        cls, input: StampedTrajectory
+        cls, input: StampedTrajectory, vel_filter_window_size: int = 1
     ) -> typing.Generator[DurationPoint, None, None]:
         """
         Generator over the trajectory.
         Yields tuples (duration in microseconds, state), state having
         a position and a velocity attribute.
         """
-        durations, positions, velocities = cls.to_duration(input)
+        durations, positions, velocities = cls.to_duration(input, vel_filter_window_size)
         for d, p, v in zip(durations, positions, velocities):
             yield d, o80.Item3dState(p, v)
         return