Add pandas_bridge.df_to_trajectory() (#656)

* Add pandas_bridge.df_to_trajectory()

Request: #647

* fix builtin name collision

* Platform-dependent int for windooooooooooowwwwwwwwwsssssss

evo/tools/pandas_bridge.py

@@ -26,15 +26,19 @@
 import numpy as np
 import pandas as pd
 
-from evo.core import trajectory, result
+from evo.core import result
+from evo.core.trajectory import PosePath3D, PoseTrajectory3D
 from evo.tools import file_interface, user
 from evo.tools.settings import SETTINGS
 
 logger = logging.getLogger(__name__)
 
+PathOrTrajectory = typing.Union[PosePath3D, PoseTrajectory3D]
+PathOrTrajectoryType = typing.Type[PathOrTrajectory]
 
-def trajectory_to_df(traj: trajectory.PosePath3D) -> pd.DataFrame:
-    if not isinstance(traj, trajectory.PosePath3D):
+
+def trajectory_to_df(traj: PosePath3D) -> pd.DataFrame:
+    if not isinstance(traj, PosePath3D):
         raise TypeError("trajectory.PosePath3D or derived required")
     poses_dict = {
         "x": traj.positions_xyz[:, 0],
@@ -45,25 +49,41 @@ def trajectory_to_df(traj: trajectory.PosePath3D) -> pd.DataFrame:
         "qy": traj.orientations_quat_wxyz[:, 2],
         "qz": traj.orientations_quat_wxyz[:, 3],
     }
-    if isinstance(traj, trajectory.PoseTrajectory3D):
+    if isinstance(traj, PoseTrajectory3D):
         index = traj.timestamps
     else:
         index = np.arange(0, traj.num_poses)
     return pd.DataFrame(data=poses_dict, index=index)
 
 
-def trajectory_stats_to_df(traj: trajectory.PosePath3D,
+def df_to_trajectory(
+    df: pd.DataFrame, as_type: typing.Optional[PathOrTrajectoryType] = None
+) -> PathOrTrajectory:
+    """
+    :param df: DataFrame created with trajectory_to_df()
+    :param as_type: Explicit output type, otherwise derived from the data.
+                    Either PosePath3D or PoseTrajectory3D.
+    """
+    quat_wxyz = df[["qw", "qx", "qy", "qz"]].to_numpy()
+    positions_xyz = df[["x", "y", "z"]].to_numpy()
+    if as_type is PosePath3D or df.index.dtype == np.int_:
+        return PosePath3D(positions_xyz, quat_wxyz)
+    timestamps = df.index.to_numpy()
+    return PoseTrajectory3D(positions_xyz, quat_wxyz, timestamps)
+
+
+def trajectory_stats_to_df(traj: PosePath3D,
                            name: typing.Optional[str] = None) -> pd.DataFrame:
-    if not isinstance(traj, trajectory.PosePath3D):
-        raise TypeError("trajectory.PosePath3D or derived required")
+    if not isinstance(traj, PosePath3D):
+        raise TypeError("PosePath3D or derived required")
     data_dict = {k: v for k, v in traj.get_infos().items() if np.isscalar(v)}
     data_dict.update(traj.get_statistics())
     index = [name] if name else ['0']
     return pd.DataFrame(data=data_dict, index=index)
 
 
 def trajectories_stats_to_df(
-        trajectories: typing.Dict[str, trajectory.PosePath3D]) -> pd.DataFrame:
+        trajectories: typing.Dict[str, PosePath3D]) -> pd.DataFrame:
     df = pd.DataFrame()
     for name, traj in trajectories.items():
         df = pd.concat((df, trajectory_stats_to_df(traj, name)))

test/test_pandas_bridge.py

@@ -0,0 +1,44 @@
+"""
+Unit test for pandas_bridge module.
+Author: Michael Grupp
+
+This file is part of evo (github.com/MichaelGrupp/evo).
+
+evo is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+evo is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with evo.  If not, see <http://www.gnu.org/licenses/>.
+"""
+
+import unittest
+
+import helpers
+from evo.core import trajectory
+from evo.tools import pandas_bridge
+
+
+class TrajectoryDataframeTest(unittest.TestCase):
+    def setUp(self) -> None:
+        self.path = helpers.fake_path(100)
+        self.trajectory = helpers.fake_trajectory(100, 0.1)
+
+    def test_back_and_forth(self):
+        for input_traj in (self.path, self.trajectory):
+            df = pandas_bridge.trajectory_to_df(input_traj)
+            output_traj = pandas_bridge.df_to_trajectory(df)
+            self.assertIsInstance(output_traj, type(input_traj))
+            self.assertEqual(input_traj, output_traj)
+
+    def test_explicit_type(self):
+        df = pandas_bridge.trajectory_to_df(self.trajectory)
+        output = pandas_bridge.df_to_trajectory(df,
+                                                as_type=trajectory.PosePath3D)
+        self.assertIsInstance(output, trajectory.PosePath3D)