✨ Enhance SignalProvider and GUI for custom arrow scaling functionality

Author: GiulioRomualdi

robot_log_visualizer/file_reader/signal_provider.py

@@ -53,7 +53,11 @@ def __init__(self, period: float):
 
         self._3d_arrows = {}
         self._3d_arrows_path_lock = QMutex()
-        self.max_arrow = 0
+
+        self._max_arrow = 0
+        self._custom_max_arrow = 0
+        self._is_custom_max_arrow_used = False
+        self._max_arrow_mutex = QMutex()
 
         self.period = period
 
@@ -213,6 +217,14 @@ def robot_state_path(self):
         locker = QMutexLocker(self.robot_state_path_lock)
         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):
@@ -308,6 +320,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 = {}
 
@@ -378,7 +398,9 @@ def register_3d_arrow(self, 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._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):

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)
@@ -93,7 +103,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):
@@ -682,9 +718,45 @@ 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,