From 5ad8e4e1584a0fcfce4b8f78d080f3a3db93a6b2 Mon Sep 17 00:00:00 2001 From: Giulio Romualdi Date: Thu, 26 Sep 2024 17:20:47 +0200 Subject: [PATCH] Add the possibility to draw 3d arrows --- .../file_reader/signal_provider.py | 30 ++++++++++++++++++ .../robot_visualizer/meshcat_provider.py | 27 ++++++++++++++++ robot_log_visualizer/ui/gui.py | 31 +++++++++++++++++-- 3 files changed, 86 insertions(+), 2 deletions(-) diff --git a/robot_log_visualizer/file_reader/signal_provider.py b/robot_log_visualizer/file_reader/signal_provider.py index 6d92ea2..a5f30dd 100644 --- a/robot_log_visualizer/file_reader/signal_provider.py +++ b/robot_log_visualizer/file_reader/signal_provider.py @@ -51,6 +51,10 @@ def __init__(self, period: float): self._3d_trajectories_path = {} self._3d_trajectories_path_lock = QMutex() + self._3d_arrows = {} + self._3d_arrows_path_lock = QMutex() + self.max_arrow = 0 + self.period = period self.data = {} @@ -319,6 +323,18 @@ def get_3d_point_at_index(self, index): return points + def get_3d_arrow_at_index(self, index): + arrows = {} + + self._3d_arrows_path_lock.lock() + + for key, value in self._3d_arrows.items(): + arrows[key] = self.get_item_from_path_at_index(value, index) + + self._3d_arrows_path_lock.unlock() + + return arrows + def get_3d_trajectory_at_index(self, index): trajectories = {} @@ -354,6 +370,20 @@ def unregister_3d_point(self, key): del self._3d_points_path[key] self._3d_points_path_lock.unlock() + def register_3d_arrow(self, key, arrow_path): + self._3d_arrows_path_lock.lock() + self._3d_arrows[key] = arrow_path + for _, value in self._3d_arrows.items(): + data, _ = self.get_item_from_path(arrow_path) + arrow = data[:, 3:] + self.max_arrow = max(np.max(np.linalg.norm(arrow, axis=1)), self.max_arrow) + self._3d_arrows_path_lock.unlock() + + def unregister_3d_arrow(self, key): + self._3d_arrows_path_lock.lock() + del self._3d_arrows[key] + self._3d_arrows_path_lock.unlock() + def register_3d_trajectory(self, key, trajectory_path): self._3d_trajectories_path_lock.lock() self._3d_trajectories_path[key] = trajectory_path diff --git a/robot_log_visualizer/robot_visualizer/meshcat_provider.py b/robot_log_visualizer/robot_visualizer/meshcat_provider.py index 3288943..e19e174 100644 --- a/robot_log_visualizer/robot_visualizer/meshcat_provider.py +++ b/robot_log_visualizer/robot_visualizer/meshcat_provider.py @@ -36,6 +36,7 @@ def __init__(self, signal_provider, period): self.env_list = ["GAZEBO_MODEL_PATH", "ROS_PACKAGE_PATH", "AMENT_PREFIX_PATH"] self._registered_3d_points = set() self._registered_3d_trajectories = dict() + self._register_3d_arrow = set() @property def state(self): @@ -56,6 +57,14 @@ def register_3d_point(self, point_path, color): radius=radius, color=color, shape_name=point_path ) + def register_3d_arrow(self, arrow_path, color): + radius = 0.02 + locker = QMutexLocker(self.meshcat_visualizer_mutex) + self._register_3d_arrow.add(arrow_path) + self._meshcat_visualizer.load_arrow( + radius=radius, color=color, shape_name=arrow_path + ) + def register_3d_trajectory(self, trajectory_path, color): locker = QMutexLocker(self.meshcat_visualizer_mutex) self._registered_3d_trajectories[trajectory_path] = (False, color) @@ -70,6 +79,11 @@ def unregister_3d_trajectory(self, trajectory_path): self._registered_3d_trajectories.pop(trajectory_path, None) self._meshcat_visualizer.delete(shape_name=trajectory_path) + def unregister_3d_arrow(self, arrow_path): + locker = QMutexLocker(self.meshcat_visualizer_mutex) + self._register_3d_arrow.remove(arrow_path) + self._meshcat_visualizer.delete(shape_name=arrow_path) + def load_model(self, considered_joints, model_name): def get_model_path_from_envs(env_list): return [ @@ -216,6 +230,19 @@ def run(self): color=self._registered_3d_trajectories[trajectory_path][1], ) + for ( + arrow_path, + arrow, + ) in self._signal_provider.get_3d_arrow_at_index(index).items(): + if arrow_path not in self._register_3d_arrow: + continue + + self._meshcat_visualizer.set_arrow_transform( + origin=arrow[0:3], + vector=arrow[3:6] / self._signal_provider.max_arrow, + shape_name=arrow_path, + ) + self.meshcat_visualizer_mutex.unlock() sleep_time = self._period - (time.time() - start) diff --git a/robot_log_visualizer/ui/gui.py b/robot_log_visualizer/ui/gui.py index fdc77a7..c99092a 100644 --- a/robot_log_visualizer/ui/gui.py +++ b/robot_log_visualizer/ui/gui.py @@ -166,6 +166,7 @@ def __init__(self, signal_provider, meshcat_provider, animation_period): self.video_items = [] self.visualized_3d_points = set() self.visualized_3d_trajectories = set() + self.visualized_3d_arrows = set() self.visualized_3d_points_colors_palette = ColorPalette() self.toolButton_on_click() @@ -731,8 +732,10 @@ def variableTreeWidget_on_right_click(self, item_position): add_3d_point_str = "Show as a 3D point" add_3d_trajectory_str = "Show as a 3D trajectory" + add_3d_arrow_str = "Show as a 3D arrow" remove_3d_point_str = "Remove the 3D point" remove_3d_trajectory_str = "Remove the 3D trajectory" + remove_3d_arrow_str = "Remove the 3D arrow" 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" @@ -784,6 +787,12 @@ def variableTreeWidget_on_right_click(self, item_position): else: menu.addAction(use_as_base_orientation_str + " (xyzw Quaternion)") + if item_size == 6: + if item_key in self.visualized_3d_arrows: + menu.addAction(remove_3d_arrow_str) + else: + menu.addAction(add_3d_arrow_str) + # show the menu action = menu.exec_(self.ui.variableTreeWidget.mapToGlobal(item_position)) if action is None: @@ -791,7 +800,11 @@ def variableTreeWidget_on_right_click(self, item_position): item_path = self.get_item_path(item) - if action.text() == add_3d_point_str or action.text() == add_3d_trajectory_str: + if ( + action.text() == add_3d_point_str + or action.text() == add_3d_trajectory_str + or action.text() == add_3d_arrow_str + ): color = next(self.visualized_3d_points_colors_palette) item.setForeground(0, QtGui.QBrush(QtGui.QColor(color.as_hex()))) @@ -802,12 +815,20 @@ def variableTreeWidget_on_right_click(self, item_position): ) self.signal_provider.register_3d_point(item_key, item_path) self.visualized_3d_points.add(item_key) - else: + elif action.text() == add_3d_trajectory_str: 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) + elif action.text() == add_3d_arrow_str: + self.meshcat_provider.register_3d_arrow( + item_key, list(color.as_normalized_rgb()) + ) + self.signal_provider.register_3d_arrow(item_key, item_path) + self.visualized_3d_arrows.add(item_key) + else: + raise ValueError("Unknown action") if action.text() == remove_3d_point_str: self.meshcat_provider.unregister_3d_point(item_key) @@ -821,6 +842,12 @@ def variableTreeWidget_on_right_click(self, item_position): self.visualized_3d_trajectories.remove(item_key) item.setForeground(0, QtGui.QBrush(QtGui.QColor(0, 0, 0))) + if action.text() == remove_3d_arrow_str: + self.meshcat_provider.unregister_3d_arrow(item_key) + self.signal_provider.unregister_3d_arrow(item_key) + self.visualized_3d_arrows.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