Merge pull request #2098 from sea-bass/meshcat-frame-viz

Add frame visualization to Meshcat visualizer

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Add inverse dynamics (`rnea`) Python and C++ example ([#2083](https://github.com/stack-of-tasks/pinocchio/pull/2083))
+- Add visualization of Frames in MeshCat viewer ([#2098](https://github.com/stack-of-tasks/pinocchio/pull/2098))
 
 ### Fixed
 

README.md

@@ -220,6 +220,7 @@ The following people have been involved in the development of **Pinocchio** and
 -   [Shubham Singh](https://github.com/shubhamsingh91) (UT Austin): second-order inverse dynamics derivatives
 -   [Stéphane Caron](https://scaron.info) (Inria): core developper
 -   [Joris Vaillant](https://github.com/jorisv) (Inria): core developer and manager of the project
+-   [Sebastian Castro](https://roboticseabass.com) (PickNik Robotics): MeshCat viewer features extension
 
 If you have participated in the development of **Pinocchio**, please add your name and contribution to this list.
 

bindings/python/pinocchio/visualize/meshcat_visualizer.py

@@ -19,6 +19,14 @@
 }
 COLOR_PRESETS = DEFAULT_COLOR_PROFILES.copy()
 
+FRAME_AXIS_POSITIONS = np.array([
+    [0, 0, 0], [1, 0, 0],
+    [0, 0, 0], [0, 1, 0],
+    [0, 0, 0], [0, 0, 1]]).astype(np.float32).T
+FRAME_AXIS_COLORS = np.array([
+    [1, 0, 0], [1, 0.6, 0],
+    [0, 1, 0], [0.6, 1, 0],
+    [0, 0, 1], [0, 0.6, 1]]).astype(np.float32).T
 
 def isMesh(geometry_object):
     """Check whether the geometry object contains a Mesh supported by MeshCat"""
@@ -477,11 +485,14 @@ def loadViewerModel(self, rootNodeName="pinocchio", color=None):
 
         # Visuals
         self.viewerVisualGroupName = self.viewerRootNodeName + "/" + "visuals"
-
         for visual in self.visual_model.geometryObjects:
             self.loadViewerGeometryObject(visual, pin.GeometryType.VISUAL, color)
         self.displayVisuals(True)
 
+        # Frames
+        self.viewerFramesGroupName = self.viewerRootNodeName + "/" + "frames"
+        self.displayFrames(False)
+
     def reload(self, new_geometry_object, geometry_type=None):
         """Reload a geometry_object given by its name and its type"""
         if geometry_type == pin.GeometryType.VISUAL:
@@ -515,6 +526,9 @@ def display(self, q=None):
         if self.display_visuals:
             self.updatePlacements(pin.GeometryType.VISUAL)
 
+        if self.display_frames:
+            self.updateFrames()
+
     def updatePlacements(self, geometry_type):
         if geometry_type == pin.GeometryType.VISUAL:
             geom_model = self.visual_model
@@ -589,6 +603,48 @@ def displayVisuals(self, visibility):
         if visibility:
             self.updatePlacements(pin.GeometryType.VISUAL)
 
+    def displayFrames(self, visibility, frame_ids=None, axis_length=0.2, axis_width=2):
+        """Set whether to display frames or not."""
+        self.display_frames = visibility
+        if visibility:
+            self.initializeFrames(frame_ids, axis_length, axis_width)
+        self.viewer[self.viewerFramesGroupName].set_property("visible", visibility)
+
+    def initializeFrames(self, frame_ids=None, axis_length=0.2, axis_width=2):
+        """Initializes the frame objects for display."""
+        import meshcat.geometry as mg
+        self.viewer[self.viewerFramesGroupName].delete()
+        self.frame_ids = []
+
+        for fid, frame in enumerate(self.model.frames):
+            if frame_ids is None or fid in frame_ids:
+                frame_viz_name = f"{self.viewerFramesGroupName}/{frame.name}"
+                self.viewer[frame_viz_name].set_object(
+                    mg.LineSegments(
+                        mg.PointsGeometry(
+                            position=axis_length * FRAME_AXIS_POSITIONS,
+                            color=FRAME_AXIS_COLORS,
+                        ),
+                        mg.LineBasicMaterial(
+                            linewidth=axis_width,
+                            vertexColors=True,
+                        ),
+                    )
+                )
+                self.frame_ids.append(fid)
+
+    def updateFrames(self):
+        """
+        Updates the frame visualizations with the latest transforms from model data.
+        """
+        pin.updateFramePlacements(self.model, self.data)
+        for fid in self.frame_ids:
+            frame_name = self.model.frames[fid].name
+            frame_viz_name = f"{self.viewerFramesGroupName}/{frame_name}"
+            self.viewer[frame_viz_name].set_transform(
+                self.data.oMf[fid].homogeneous
+            )
+
     def drawFrameVelocities(
         self, frame_id, v_scale=0.2, color=FRAME_VEL_COLOR
     ):  # pylint: disable=arguments-differ