Merge pull request #76 from ami-iit/draw_3d_trajectory

✨ Add the possibility to draw 3d trajectory in meshcat-visualizer

Author: GiulioRomualdi

robot_log_visualizer/file_reader/signal_provider.py

@@ -48,6 +48,9 @@ def __init__(self, period: float):
         self._3d_points_path = {}
         self._3d_points_path_lock = QMutex()
 
+        self._3d_trajectories_path = {}
+        self._3d_trajectories_path_lock = QMutex()
+
         self.period = period
 
         self.data = {}
@@ -64,6 +67,8 @@ def __init__(self, period: float):
 
         self._current_time = 0
 
+        self.trajectory_span = 200
+
     def __populate_text_logging_data(self, file_object):
         data = {}
         for key, value in file_object.items():
@@ -241,12 +246,18 @@ def get_item_from_path(self, path, default_path=None):
 
         return data["data"], data["timestamps"]
 
-    def get_item_from_path_at_index(self, path, index, default_path=None):
+    def get_item_from_path_at_index(self, path, index, default_path=None, neighbor=0):
         data, timestamps = self.get_item_from_path(path, default_path)
         if data is None:
             return None
         closest_index = np.argmin(np.abs(timestamps - self.timestamps[index]))
-        return data[closest_index, :]
+
+        if neighbor == 0:
+            return data[closest_index, :]
+
+        initial_index = max(0, closest_index - neighbor)
+        end_index = min(len(timestamps), closest_index + neighbor + 1)
+        return data[initial_index:end_index, :]
 
     def get_robot_state_at_index(self, index):
         robot_state = {}
@@ -308,6 +319,31 @@ def get_3d_point_at_index(self, index):
 
         return points
 
+    def get_3d_trajectory_at_index(self, index):
+        trajectories = {}
+
+        self._3d_trajectories_path_lock.lock()
+
+        for key, value in self._3d_trajectories_path.items():
+            trajectories[key] = self.get_item_from_path_at_index(
+                value, index, neighbor=self.trajectory_span
+            )
+            # force the size of the points to be 3 if less than 3 we assume that the point is a 2d point and we add a 0 as z coordinate
+            if trajectories[key].shape[1] < 3:
+                trajectories[key] = np.concatenate(
+                    (
+                        trajectories[key],
+                        np.zeros(
+                            (trajectories[key].shape[0], 3 - trajectories[key].shape[1])
+                        ),
+                    ),
+                    axis=1,
+                )
+
+        self._3d_trajectories_path_lock.unlock()
+
+        return trajectories
+
     def register_3d_point(self, key, points_path):
         self._3d_points_path_lock.lock()
         self._3d_points_path[key] = points_path
@@ -318,6 +354,16 @@ def unregister_3d_point(self, key):
         del self._3d_points_path[key]
         self._3d_points_path_lock.unlock()
 
+    def register_3d_trajectory(self, key, trajectory_path):
+        self._3d_trajectories_path_lock.lock()
+        self._3d_trajectories_path[key] = trajectory_path
+        self._3d_trajectories_path_lock.unlock()
+
+    def unregister_3d_trajectory(self, key):
+        self._3d_trajectories_path_lock.lock()
+        del self._3d_trajectories_path[key]
+        self._3d_trajectories_path_lock.unlock()
+
     def run(self):
         while True:
             start = time.time()

robot_log_visualizer/robot_visualizer/meshcat_provider.py

@@ -35,6 +35,7 @@ def __init__(self, signal_provider, period):
         self.custom_package_dir = ""
         self.env_list = ["GAZEBO_MODEL_PATH", "ROS_PACKAGE_PATH", "AMENT_PREFIX_PATH"]
         self._registered_3d_points = set()
+        self._registered_3d_trajectories = dict()
 
     @property
     def state(self):
@@ -55,11 +56,20 @@ def register_3d_point(self, point_path, color):
             radius=radius, color=color, shape_name=point_path
         )
 
+    def register_3d_trajectory(self, trajectory_path, color):
+        locker = QMutexLocker(self.meshcat_visualizer_mutex)
+        self._registered_3d_trajectories[trajectory_path] = (False, color)
+
     def unregister_3d_point(self, point_path):
         locker = QMutexLocker(self.meshcat_visualizer_mutex)
         self._registered_3d_points.remove(point_path)
         self._meshcat_visualizer.delete(shape_name=point_path)
 
+    def unregister_3d_trajectory(self, trajectory_path):
+        locker = QMutexLocker(self.meshcat_visualizer_mutex)
+        self._registered_3d_trajectories.pop(trajectory_path, None)
+        self._meshcat_visualizer.delete(shape_name=trajectory_path)
+
     def load_model(self, considered_joints, model_name):
         def get_model_path_from_envs(env_list):
             return [
@@ -162,11 +172,9 @@ def run(self):
         while True:
             start = time.time()
 
+            index = self._signal_provider.index
             if self.state == PeriodicThreadState.running and self._is_model_loaded:
-                robot_state = self._signal_provider.get_robot_state_at_index(
-                    self._signal_provider.index
-                )
-
+                robot_state = self._signal_provider.get_robot_state_at_index(index)
                 self.meshcat_visualizer_mutex.lock()
                 # These are the robot measured joint positions in radians
                 self._meshcat_visualizer.set_multibody_system_state(
@@ -177,7 +185,7 @@ def run(self):
                 )
 
                 for points_path, points in self._signal_provider.get_3d_point_at_index(
-                    self._signal_provider.index
+                    index
                 ).items():
                     if points_path not in self._registered_3d_points:
                         continue
@@ -186,6 +194,28 @@ def run(self):
                         position=points, rotation=identity, shape_name=points_path
                     )
 
+                for (
+                    trajectory_path,
+                    trajectory,
+                ) in self._signal_provider.get_3d_trajectory_at_index(index).items():
+                    if trajectory_path not in self._registered_3d_trajectories.keys():
+                        continue
+
+                    if self._registered_3d_trajectories[trajectory_path][0]:
+                        self._meshcat_visualizer.delete(shape_name=trajectory_path)
+                    else:
+                        self._registered_3d_trajectories[trajectory_path] = (
+                            True,
+                            self._registered_3d_trajectories[trajectory_path][1],
+                        )
+
+                    self._meshcat_visualizer.load_line(
+                        vertices=trajectory.T,
+                        linewidth=5.0,
+                        shape_name=trajectory_path,
+                        color=self._registered_3d_trajectories[trajectory_path][1],
+                    )
+
                 self.meshcat_visualizer_mutex.unlock()
 
             sleep_time = self._period - (time.time() - start)

robot_log_visualizer/ui/gui.py

@@ -165,6 +165,7 @@ def __init__(self, signal_provider, meshcat_provider, animation_period):
         self.plot_items = []
         self.video_items = []
         self.visualized_3d_points = set()
+        self.visualized_3d_trajectories = set()
         self.visualized_3d_points_colors_palette = ColorPalette()
 
         self.toolButton_on_click()
@@ -722,7 +723,9 @@ def variableTreeWidget_on_right_click(self, item_position):
         menu = QtWidgets.QMenu()
 
         add_3d_point_str = "Show as a 3D point"
+        add_3d_trajectory_str = "Show as a 3D trajectory"
         remove_3d_point_str = "Remove the 3D point"
+        remove_3d_trajectory_str = "Remove the 3D trajectory"
         use_as_base_position_str = "Use as base position"
         use_as_base_orientation_str = "Use as base orientation"
         dont_use_as_base_position_str = "Don't use as base position"
@@ -734,8 +737,14 @@ def variableTreeWidget_on_right_click(self, item_position):
         if item_size == 2:
             if item_key in self.visualized_3d_points:
                 menu.addAction(remove_3d_point_str)
-            else:
+            if item_key in self.visualized_3d_trajectories:
+                menu.addAction(remove_3d_trajectory_str)
+            if (
+                item_key not in self.visualized_3d_points
+                and item_key not in self.visualized_3d_trajectories
+            ):
                 menu.addAction(add_3d_point_str)
+                menu.addAction(add_3d_trajectory_str)
 
         # in this case we can use the item as base position, base orientation or 3d point
         if item_size == 3:
@@ -753,8 +762,14 @@ def variableTreeWidget_on_right_click(self, item_position):
 
             if item_key in self.visualized_3d_points:
                 menu.addAction(remove_3d_point_str)
-            else:
+            if item_key in self.visualized_3d_trajectories:
+                menu.addAction(remove_3d_trajectory_str)
+            if (
+                item_key not in self.visualized_3d_points
+                and item_key not in self.visualized_3d_trajectories
+            ):
                 menu.addAction(add_3d_point_str)
+                menu.addAction(add_3d_trajectory_str)
 
         if item_size == 4:
             if item_path == self.robot_state_path.base_orientation_path:
@@ -769,22 +784,36 @@ def variableTreeWidget_on_right_click(self, item_position):
 
         item_path = self.get_item_path(item)
 
-        if action.text() == add_3d_point_str:
+        if action.text() == add_3d_point_str or action.text() == add_3d_trajectory_str:
             color = next(self.visualized_3d_points_colors_palette)
 
             item.setForeground(0, QtGui.QBrush(QtGui.QColor(color.as_hex())))
-            self.meshcat_provider.register_3d_point(
-                item_key, list(color.as_normalized_rgb())
-            )
-            self.signal_provider.register_3d_point(item_key, item_path)
-            self.visualized_3d_points.add(item_key)
+
+            if action.text() == add_3d_point_str:
+                self.meshcat_provider.register_3d_point(
+                    item_key, list(color.as_normalized_rgb())
+                )
+                self.signal_provider.register_3d_point(item_key, item_path)
+                self.visualized_3d_points.add(item_key)
+            else:
+                self.meshcat_provider.register_3d_trajectory(
+                    item_key, list(color.as_normalized_rgb())
+                )
+                self.signal_provider.register_3d_trajectory(item_key, item_path)
+                self.visualized_3d_trajectories.add(item_key)
 
         if action.text() == remove_3d_point_str:
             self.meshcat_provider.unregister_3d_point(item_key)
             self.signal_provider.unregister_3d_point(item_key)
             self.visualized_3d_points.remove(item_key)
             item.setForeground(0, QtGui.QBrush(QtGui.QColor(0, 0, 0)))
 
+        if action.text() == remove_3d_trajectory_str:
+            self.meshcat_provider.unregister_3d_trajectory(item_key)
+            self.signal_provider.unregister_3d_trajectory(item_key)
+            self.visualized_3d_trajectories.remove(item_key)
+            item.setForeground(0, QtGui.QBrush(QtGui.QColor(0, 0, 0)))
+
         if (
             use_as_base_orientation_str in action.text()
             or action.text() == use_as_base_position_str

setup.cfg

@@ -40,7 +40,7 @@ classifiers =
 packages = find:
 python_requires = >=3.8
 install_requires =
-   idyntree >= 8.0.0
+   idyntree >= 10.2.0
    meshcat
    numpy
    PyQt5