diff --git a/robot_log_visualizer/file_reader/signal_provider.py b/robot_log_visualizer/file_reader/signal_provider.py index dc49e9e..6d92ea2 100644 --- a/robot_log_visualizer/file_reader/signal_provider.py +++ b/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() diff --git a/robot_log_visualizer/robot_visualizer/meshcat_provider.py b/robot_log_visualizer/robot_visualizer/meshcat_provider.py index b39b25d..3288943 100644 --- a/robot_log_visualizer/robot_visualizer/meshcat_provider.py +++ b/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) diff --git a/robot_log_visualizer/ui/gui.py b/robot_log_visualizer/ui/gui.py index 0204407..e3428d4 100644 --- a/robot_log_visualizer/ui/gui.py +++ b/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,15 +784,23 @@ 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) @@ -785,6 +808,12 @@ def variableTreeWidget_on_right_click(self, item_position): 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 diff --git a/setup.cfg b/setup.cfg index 3c8fb7e..2684018 100644 --- a/setup.cfg +++ b/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