Separate initialization and update of frames

stack-of-tasks · Nov 25, 2023 · 1186f62 · 1186f62
1 parent eb853d0
commit 1186f62
Showing 1 changed file with 43 additions and 32 deletions.
diff --git a/bindings/python/pinocchio/visualize/meshcat_visualizer.py b/bindings/python/pinocchio/visualize/meshcat_visualizer.py
@@ -485,14 +485,14 @@ def loadViewerModel(self, rootNodeName="pinocchio", color=None):
 
         # Visuals
         self.viewerVisualGroupName = self.viewerRootNodeName + "/" + "visuals"
-
-        # Frames
-        self.viewerFramesGroupName = self.viewerRootNodeName + "/" + "frames"
-
         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:
@@ -526,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
@@ -600,38 +603,46 @@ def displayVisuals(self, visibility):
         if visibility:
             self.updatePlacements(pin.GeometryType.VISUAL)
 
-    def drawFrame(
-        self, frame_id, axis_length=0.2, axis_width=2
-    ):
+    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
-        frame_name = self.model.frames[frame_id].name
-        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.viewer[frame_viz_name].set_transform(
-            self.data.oMf[frame_id].homogeneous
-        )
+        self.viewer[self.viewerFramesGroupName].delete()
+        self.frame_ids = []
 
-    def drawFrames(
-        self, frame_ids=None, axis_length=0.2, axis_width=2
-    ):
-        for fid, _ in enumerate(self.model.frames):
+        for fid, frame in enumerate(self.model.frames):
             if frame_ids is None or fid in frame_ids:
-                self.drawFrame(
-                    fid,
-                    axis_length=axis_length,
-                    axis_width=axis_width,
+                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.
+        """
+        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