Add the possibility to draw 3d arrows

robot_log_visualizer/file_reader/signal_provider.py

@@ -51,6 +51,14 @@ 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._custom_max_arrow = 0
+        self._is_custom_max_arrow_used = False
+        self._max_arrow_mutex = QMutex()
+
         self.period = period
 
         self.data = {}
@@ -88,14 +96,16 @@ def __populate_text_logging_data(self, file_object):
                     # If len(value[text[0]].shape) == 2 then the text contains a string, otherwise it is empty
                     # We need to manually check the shape to handle the case in which the text is empty
                     data[key]["data"] = [
-                        TextLoggingMsg(
-                            text="".join(chr(c[0]) for c in value[text[0]]),
-                            level="".join(chr(c[0]) for c in value[level[0]]),
-                        )
-                        if len(value[text[0]].shape) == 2
-                        else TextLoggingMsg(
-                            text="",
-                            level="".join(chr(c[0]) for c in value[level[0]]),
+                        (
+                            TextLoggingMsg(
+                                text="".join(chr(c[0]) for c in value[text[0]]),
+                                level="".join(chr(c[0]) for c in value[level[0]]),
+                            )
+                            if len(value[text[0]].shape) == 2
+                            else TextLoggingMsg(
+                                text="",
+                                level="".join(chr(c[0]) for c in value[level[0]]),
+                            )
                         )
                         for text, level in zip(text_ref, level_ref)
                     ]
@@ -210,6 +220,14 @@ def robot_state_path(self):
         value = self._robot_state_path
         return value
 
+    @property
+    def max_arrow(self):
+        locker = QMutexLocker(self._max_arrow_mutex)
+        if self._is_custom_max_arrow_used:
+            return self._custom_max_arrow
+        else:
+            return self._max_arrow
+
     @robot_state_path.setter
     def robot_state_path(self, robot_state_path):
         locker = QMutexLocker(self.robot_state_path_lock)
@@ -304,6 +322,14 @@ def get_robot_state_at_index(self, index):
 
         return robot_state
 
+    def set_custom_max_arrow(self, use_custom_max_arrow: bool, max_arrow: float):
+        _ = QMutexLocker(self._max_arrow_mutex)
+        self._is_custom_max_arrow_used = use_custom_max_arrow
+        if use_custom_max_arrow:
+            self._custom_max_arrow = max_arrow
+        else:
+            self._custom_max_arrow = 0
+
     def get_3d_point_at_index(self, index):
         points = {}
 
@@ -321,6 +347,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 = {}
 
@@ -356,6 +394,24 @@ 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_mutex.lock()
+            self._max_arrow = max(
+                np.max(np.linalg.norm(arrow, axis=1)), self._max_arrow
+            )
+            self._max_arrow_mutex.unlock()
+        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

robot_log_visualizer/robot_visualizer/meshcat_provider.py

@@ -37,6 +37,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()
 
         self.frame_T_base = np.eye(4)
 
@@ -59,6 +60,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)
@@ -73,6 +82,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, base_frame=None):
         def get_model_path_from_envs(env_list):
             return [
@@ -269,6 +283,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)

robot_log_visualizer/ui/gui.py

@@ -71,6 +71,16 @@ def __init__(self, meshcat_provider, parent=None, dataset_loaded=False):
         self.ui.robotModelToolButton.clicked.connect(self.open_urdf_file)
         self.ui.packageDirToolButton.clicked.connect(self.open_package_directory)
 
+        # Force the arrowScaling_lineEdit to be a positive float
+        self.ui.arrowScaling_lineEdit.setValidator(QtGui.QDoubleValidator(0, 100, 2))
+
+        # connect the arrowScaling_checkBox to the handle_arrow_scaling method
+        self.ui.arrowScaling_checkBox.toggled.connect(self.handle_arrow_scaling)
+
+        self.clicked_button = None
+        self.std_button = None
+        self.ui.buttonBox.clicked.connect(self.buttonBox_on_click)
+
         if dataset_loaded:
             frames = meshcat_provider.robot_frames()
             self.ui.frameNameComboBox.addItems(frames)
@@ -94,6 +104,33 @@ def get_urdf_path(self):
     def get_package_directory(self):
         return self.ui.packageDirLineEdit.text()
 
+    def buttonBox_on_click(self, button):
+        self.clicked_button = button
+
+        self.std_button = self.ui.buttonBox.standardButton(button)
+
+    def get_clicked_button_role(self):
+        if self.clicked_button is not None:
+            return self.ui.buttonBox.buttonRole(self.clicked_button)
+        return None
+
+    def get_clicked_button_text(self):
+        if self.clicked_button is not None:
+            return self.clicked_button.text()
+        return None
+
+    def get_clicked_standard_button(self):
+        return self.std_button
+
+    def handle_arrow_scaling(self):
+        # if arrowScaling_checkBox is checked the lineEdit must be disabled else it must be enabled
+        if self.ui.arrowScaling_checkBox.isChecked():
+            self.ui.arrowScaling_lineEdit.setText("")
+            self.ui.arrowScaling_lineEdit.setEnabled(False)
+        else:
+            self.ui.arrowScaling_lineEdit.setText("")
+            self.ui.arrowScaling_lineEdit.setEnabled(True)
+
 
 class About(QtWidgets.QMainWindow):
     def __init__(self):
@@ -172,6 +209,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()
@@ -681,9 +719,47 @@ def open_set_robot_model(self):
         )
         outcome = dlg.exec()
         if outcome == QDialog.Accepted:
-            if not self.dataset_loaded:
-                self.meshcat_provider.model_path = dlg.get_urdf_path()
-                self.meshcat_provider.custom_package_dir = dlg.get_package_directory()
+
+            # check which button was clicked
+            button_role = dlg.get_clicked_button_role()
+            button_text = dlg.get_clicked_button_text()
+            std_button = dlg.get_clicked_standard_button()
+
+            if std_button == QtWidgets.QDialogButtonBox.SaveAll:
+                if not self.dataset_loaded:
+                    self.meshcat_provider.model_path = dlg.get_urdf_path()
+                    self.meshcat_provider.custom_package_dir = (
+                        dlg.get_package_directory()
+                    )
+
+                arrow_scaling_value = dlg.ui.arrowScaling_lineEdit.text()
+                if not arrow_scaling_value:
+                    arrow_scaling_value = "1.0"
+                else:
+                    arrow_scaling_value = float(arrow_scaling_value)
+                self.signal_provider.set_custom_max_arrow(
+                    not dlg.ui.arrowScaling_checkBox.isChecked(), arrow_scaling_value
+                )
+            if std_button == QtWidgets.QDialogButtonBox.Save:
+                # we need to check which tab is selected in the dlg
+                if dlg.ui.tabWidget.currentIndex() == 0:
+                    if not self.dataset_loaded:
+                        self.meshcat_provider.model_path = dlg.get_urdf_path()
+                        self.meshcat_provider.custom_package_dir = (
+                            dlg.get_package_directory()
+                        )
+                else:
+                    arrow_scaling_value = dlg.ui.arrowScaling_lineEdit.text()
+                    # if it is empty we set it to 1.0
+                    if not arrow_scaling_value:
+                        arrow_scaling_value = "1.0"
+                    else:
+                        arrow_scaling_value = float(arrow_scaling_value)
+                    self.signal_provider.set_custom_max_arrow(
+                        not dlg.ui.arrowScaling_checkBox.isChecked(),
+                        arrow_scaling_value,
+                    )
+
             else:
                 self.meshcat_provider.load_model(
                     self.signal_provider.joints_name,
@@ -743,8 +819,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"
@@ -796,14 +874,24 @@ 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:
             return
 
         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())))
@@ -814,12 +902,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)
@@ -833,6 +929,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