Skip to content

Commit

Permalink
Add the possibility to change the baseframe in the visualizer
Browse files Browse the repository at this point in the history
  • Loading branch information
@GiulioRomualdi
GiulioRomualdi committed Feb 7, 2025
1 parent 3f26b66 commit bcd7af0
Show file tree
Hide file tree
Showing 4 changed files with 142 additions and 60 deletions.
4 changes: 3 additions & 1 deletion robot_log_visualizer/file_reader/signal_provider.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,9 @@ def __populate_numerical_data(self, file_object):
if "data" in value.keys():
data[key] = {}
data[key]["data"] = np.atleast_1d(np.squeeze(np.array(value["data"])))
data[key]["timestamps"] = np.atleast_1d(np.squeeze(np.array(value["timestamps"])))
data[key]["timestamps"] = np.atleast_1d(
np.squeeze(np.array(value["timestamps"]))
)

# if the initial or end time has been updated we can also update the entire timestamps dataset
if data[key]["timestamps"][0] < self.initial_time:
Expand Down
71 changes: 58 additions & 13 deletions robot_log_visualizer/robot_visualizer/meshcat_provider.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,12 +31,15 @@ def __init__(self, signal_provider, period):
self._is_model_loaded = False
self._signal_provider = signal_provider

self.custom_model_path = ""
self.model_path = ""
self.custom_package_dir = ""
self.base_frame = ""
self.env_list = ["GAZEBO_MODEL_PATH", "ROS_PACKAGE_PATH", "AMENT_PREFIX_PATH"]
self._registered_3d_points = set()
self._registered_3d_trajectories = dict()

self.frame_T_base = np.eye(4)

@property
def state(self):
locker = QMutexLocker(self.state_lock)
Expand Down Expand Up @@ -70,7 +73,7 @@ def unregister_3d_trajectory(self, trajectory_path):
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 load_model(self, considered_joints, model_name, base_frame=None):
def get_model_path_from_envs(env_list):
return [
Path(f) if (env != "AMENT_PREFIX_PATH") else Path(f) / "share"
Expand Down Expand Up @@ -101,23 +104,21 @@ def find_model_joints(model_name, considered_joints):

return model_joints_index

self._is_model_loaded = False

# Load the model
model_loader = idyn.ModelLoader()

self.model_joints_index = []
# In this case the user specify the model path
if self.custom_model_path:
if self.custom_package_dir:
self.model_joints_index = find_model_joints(
self.custom_model_path, considered_joints
self.model_path, considered_joints
)
considered_model_joints = [
considered_joints[i] for i in self.model_joints_index
]

model_loader.loadReducedModelFromFile(
self.custom_model_path,
self.model_path,
considered_model_joints,
"urdf",
[self.custom_package_dir],
Expand All @@ -138,34 +139,71 @@ def find_model_joints(model_name, considered_joints):

if model_filenames:
model_found_in_env_folders = True
self.custom_model_path = str(model_filenames[0])
self.model_path = str(model_filenames[0])
break

# If the model is not found we exit
if not model_found_in_env_folders:
return False

self.model_joints_index = find_model_joints(
self.custom_model_path, considered_joints
self.model_path, considered_joints
)
considered_model_joints = [
considered_joints[i] for i in self.model_joints_index
]
model_loader.loadReducedModelFromFile(
self.custom_model_path, considered_model_joints
self.model_path, considered_model_joints
)

if not model_loader.isValid():
return False

self.meshcat_visualizer_mutex.lock()

# if the model is already loaded we remove it
if self._is_model_loaded:
self._meshcat_visualizer.delete(shape_name="robot")
self._is_model_loaded = False

model = model_loader.model()
if base_frame is None:
self.frame_T_base = np.eye(4)
self.base_frame = model.getFrameName(model.getDefaultBaseLink())
link_frame = None
else:
base_frame_index = model.getFrameIndex(base_frame)
frame_T_base = model.getFrameTransform(base_frame_index)
frame_T_base = frame_T_base.inverse().asHomogeneousTransform()
self.frame_T_base = frame_T_base.toNumPy()

link_frame = model.getFrameName(model.getFrameLink(base_frame_index))
self.base_frame = base_frame

self._meshcat_visualizer.load_model(
model_loader.model(), model_name="robot", color=0.8
model, model_name="robot", color=0.8, base_frame=link_frame
)

self._is_model_loaded = True

self.meshcat_visualizer_mutex.unlock()

return True

def robot_frames(self):
frames = []
if not self._is_model_loaded:
return frames

model = self._meshcat_visualizer.model["robot"]
frames_number = model.getNrOfFrames()

for i in range(frames_number):
frame = model.getFrameName(i)
frames.append(frame)

return frames

def run(self):
identity = np.eye(3)

Expand All @@ -177,9 +215,16 @@ def run(self):
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
# we need to compute the base transform as
# I_T_B = I_T_F * frame_T_Base
# I_R_B = I_R_F * frame_R_Base
# I_p_B = I_p_F + I_R_F * frame_p_Base
self._meshcat_visualizer.set_multibody_system_state(
base_position=robot_state["base_position"],
base_rotation=robot_state["base_orientation"],
base_position=robot_state["base_orientation"]
@ self.frame_T_base[:3, 3]
+ robot_state["base_position"],
base_rotation=robot_state["base_orientation"]
@ self.frame_T_base[:3, :3],
joint_value=robot_state["joints_position"][self.model_joints_index],
model_name="robot",
)
Expand Down
36 changes: 24 additions & 12 deletions robot_log_visualizer/ui/gui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,26 +50,32 @@


class SetRobotModelDialog(QtWidgets.QDialog):
def __init__(
self, parent=None, model_path=None, package_dir=None, model_modificable=True
):
def __init__(self, meshcat_provider, parent=None, dataset_loaded=False):
# call QMainWindow constructor
super().__init__(parent)
self.ui = Ui_setRobotModelDialog()
self.ui.setupUi(self)

model_path = meshcat_provider.model_path
package_dir = meshcat_provider.custom_package_dir

if model_path:
self.ui.robotModelLineEdit.setText(model_path)

if package_dir:
self.ui.packageDirLineEdit.setText(package_dir)

self.ui.robotModelLineEdit.setEnabled(model_modificable)
self.ui.packageDirLineEdit.setEnabled(model_modificable)
self.ui.robotModelLineEdit.setEnabled(not dataset_loaded)
self.ui.packageDirLineEdit.setEnabled(not dataset_loaded)

self.ui.robotModelToolButton.clicked.connect(self.open_urdf_file)
self.ui.packageDirToolButton.clicked.connect(self.open_package_directory)

if dataset_loaded:
frames = meshcat_provider.robot_frames()
self.ui.frameNameComboBox.addItems(frames)
self.ui.frameNameComboBox.setCurrentText(meshcat_provider.base_frame)

def open_urdf_file(self):
file_name, _ = QFileDialog.getOpenFileName(
self, "Open urdf file", ".", filter="*.urdf"
Expand Down Expand Up @@ -589,8 +595,8 @@ def __load_mat_file(self, file_name):
):
# if not loaded we print an error but we continue
msg = "Unable to load the model: "
if self.meshcat_provider.custom_model_path:
msg = msg + self.meshcat_provider.custom_model_path
if self.meshcat_provider.model_path:
msg = msg + self.meshcat_provider.model_path
else:
msg = msg + self.signal_provider.robot_name

Expand Down Expand Up @@ -669,15 +675,21 @@ def open_about(self):

def open_set_robot_model(self):
dlg = SetRobotModelDialog(
self.meshcat_provider,
self,
self.meshcat_provider.custom_model_path,
self.meshcat_provider.custom_package_dir,
not self.dataset_loaded,
self.dataset_loaded,
)
outcome = dlg.exec()
if outcome == QDialog.Accepted:
self.meshcat_provider.custom_model_path = dlg.get_urdf_path()
self.meshcat_provider.custom_package_dir = dlg.get_package_directory()
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:
self.meshcat_provider.load_model(
self.signal_provider.joints_name,
self.signal_provider.robot_name,
base_frame=dlg.ui.frameNameComboBox.currentText(),
)

def dropEvent(self, event):
if len(event.mimeData().urls()) != 1:
Expand Down
91 changes: 57 additions & 34 deletions robot_log_visualizer/ui/misc/set_robot_model.ui
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,39 +6,20 @@
<rect>
<x>0</x>
<y>0</y>
<width>365</width>
<height>217</height>
<width>711</width>
<height>363</height>
</rect>
</property>
<property name="windowTitle">
<string>Dialog</string>
</property>
<layout class="QGridLayout" name="gridLayout">
<item row="3" column="0">
<widget class="QDialogButtonBox" name="buttonBox">
<property name="orientation">
<enum>Qt::Horizontal</enum>
</property>
<property name="standardButtons">
<set>QDialogButtonBox::Abort|QDialogButtonBox::Save</set>
</property>
</widget>
</item>
<item row="0" column="0">
<layout class="QFormLayout" name="formLayout">
<item row="0" column="0" colspan="2">
<widget class="QFrame" name="frame">
<property name="frameShape">
<enum>QFrame::StyledPanel</enum>
</property>
<layout class="QGridLayout" name="gridLayout_2">
<item row="1" column="0">
<widget class="QLineEdit" name="robotModelLineEdit">
<property name="font">
<font>
<family>Ubuntu Mono</family>
</font>
</property>
</widget>
</item>
<item row="0" column="0">
<widget class="QLabel" name="label">
<property name="text">
Expand All @@ -53,42 +34,84 @@
</property>
</widget>
</item>
<item row="1" column="0">
<widget class="QLineEdit" name="robotModelLineEdit">
<property name="font">
<font>
<family>Ubuntu Mono</family>
</font>
</property>
</widget>
</item>
</layout>
</widget>
</item>
<item row="1" column="0">
<item row="1" column="0" colspan="2">
<widget class="QFrame" name="frame_2">
<property name="frameShape">
<enum>QFrame::StyledPanel</enum>
</property>
<layout class="QGridLayout" name="gridLayout_3">
<item row="1" column="0">
<widget class="QLineEdit" name="packageDirLineEdit">
<property name="font">
<font>
<family>Ubuntu Mono</family>
</font>
</property>
</widget>
</item>
<item row="1" column="1">
<widget class="QToolButton" name="packageDirToolButton">
<property name="text">
<string>...</string>
</property>
</widget>
</item>
<item row="0" column="0">
<widget class="QLabel" name="label_3">
<property name="text">
<string>Package Directory</string>
</property>
</widget>
</item>
<item row="1" column="1">
<widget class="QToolButton" name="packageDirToolButton">
</layout>
</widget>
</item>
<item row="2" column="0" colspan="2">
<widget class="QFrame" name="frame_3">
<property name="frameShape">
<enum>QFrame::StyledPanel</enum>
</property>
<layout class="QFormLayout" name="formLayout_2">
<item row="0" column="0">
<widget class="QLabel" name="label_5">
<property name="text">
<string>...</string>
<string>Base Frame</string>
</property>
</widget>
</item>
<item row="1" column="0">
<widget class="QLineEdit" name="packageDirLineEdit">
<property name="font">
<font>
<family>Ubuntu Mono</family>
</font>
<item row="0" column="1">
<widget class="QComboBox" name="frameNameComboBox">
<property name="maxVisibleItems">
<number>5</number>
</property>
</widget>
</item>
</layout>
</widget>
</item>
<item row="2" column="0">
<item row="4" column="0" colspan="2">
<widget class="QDialogButtonBox" name="buttonBox">
<property name="orientation">
<enum>Qt::Horizontal</enum>
</property>
<property name="standardButtons">
<set>QDialogButtonBox::Abort|QDialogButtonBox::Save</set>
</property>
</widget>
</item>
<item row="3" column="1">
<spacer name="verticalSpacer">
<property name="orientation">
<enum>Qt::Vertical</enum>
Expand Down

0 comments on commit bcd7af0

Please sign in to comment.