Add color support for robot meshes in Viser

bindings/python/pinocchio/visualize/viser_visualizer.py

@@ -1,4 +1,5 @@
 import time
+import os
 
 try:
     import hppfcl
@@ -7,12 +8,15 @@
 
 import numpy as np
 
+from trimesh.exchange import dae
 from .. import pinocchio_pywrap_default as pin
 from . import BaseVisualizer
 
 try:
     import trimesh  # Required by viser
     import viser
+    import collada
+
 except ImportError:
     import_viser_succeed = False
 else:
@@ -214,20 +218,106 @@ def loadViewerGeometryObject(self, geometry_object, prefix="", color=None):
 
         self.frames[name] = frame
 
-    def _add_mesh_from_path(self, name, mesh_path, color):
+    def _add_mesh_from_path(self, name, mesh_path, color=None):
         """
         Load a mesh from a file.
         """
+        ext = os.path.splitext(mesh_path)[1].lower()
+
+        if ext == ".dae":
+            try:
+                mesh_collada = collada.Collada(mesh_path)
+
+                if len(mesh_collada.effects) < len(mesh_collada.geometries):
+                    mesh = trimesh.load_mesh(mesh_path)
+                    if color is None:
+                        return self.viewer.scene.add_mesh_trimesh(name, mesh)
+                    else:
+                        return self.viewer.scene.add_mesh_simple(
+                            name, mesh.vertices, mesh.faces,
+                            color=color[:3], opacity=color[3]
+                        )
+
+                frames = []
+                for i, (geometry, effect) in enumerate(
+                    zip(mesh_collada.geometries, mesh_collada.effects)
+                ):
+                    try:
+                        vertices = geometry.primitives[0].sources["VERTEX"][0][4].data
+                        indices = geometry.primitives[0].indices
+
+                        if indices.ndim == 3:
+                            faces = indices[:, :, 0]
+                        elif indices.ndim == 2:
+                            faces = indices
+                        elif indices.ndim == 1:
+                            faces = indices.reshape(-1, 3)
+                        else:
+                            faces = indices.reshape(-1, 3)
+
+                        if faces.shape[-1] != 3:
+                            if faces.size % 3 == 0:
+                                faces = faces.reshape(-1, 3)
+
+                        mesh_color = getattr(effect, "diffuse", None)
+
+                        if mesh_color is not None:
+                            frame = self.viewer.scene.add_mesh_simple(
+                                f"{name}_{i}",
+                                vertices,
+                                faces,
+                                color=mesh_color[:3],
+                                opacity=mesh_color[3],
+                            )
+                        elif color is not None:
+                            frame = self.viewer.scene.add_mesh_simple(
+                                f"{name}_{i}",
+                                vertices,
+                                faces,
+                                color=color[:3],
+                                opacity=color[3],
+                            )
+                        else:
+                            mesh = trimesh.load_mesh(mesh_path)
+                            frame = self.viewer.scene.add_mesh_trimesh(f"{name}_{i}", mesh)
+
+                        if frame is None:
+                            frame = self.viewer.scene.add_box(
+                                f"{name}_{i}_placeholder",
+                                (0.001, 0.001, 0.001),
+                                color=(200, 200, 200)
+                            )
+
+                        frames.append(frame)
+
+                    except Exception:
+                        try:
+                            mesh = trimesh.load_mesh(mesh_path)
+                            frame = self.viewer.scene.add_mesh_trimesh(f"{name}_{i}", mesh)
+                            if frame is not None:
+                                frames.append(frame)
+                        except Exception:
+                            pass
+
+                return frames[0] if frames else None
+
+            except Exception:
+                mesh = trimesh.load_mesh(mesh_path)
+                if color is None:
+                    return self.viewer.scene.add_mesh_trimesh(name, mesh)
+                else:
+                    return self.viewer.scene.add_mesh_simple(
+                        name, mesh.vertices, mesh.faces,
+                        color=color[:3], opacity=color[3]
+                    )
+
         mesh = trimesh.load_mesh(mesh_path)
         if color is None:
             return self.viewer.scene.add_mesh_trimesh(name, mesh)
         else:
             return self.viewer.scene.add_mesh_simple(
-                name,
-                mesh.vertices,
-                mesh.faces,
-                color=color[:3],
-                opacity=color[3],
+                name, mesh.vertices, mesh.faces,
+                color=color[:3], opacity=color[3]
             )
 
     def _add_mesh_from_convex(self, name, geom, color):

examples/gepetto-viewer.py

@@ -11,8 +11,8 @@
 # Load the URDF model.
 model_path = Path(os.environ.get("EXAMPLE_ROBOT_DATA_MODEL_DIR"))
 mesh_dir = model_path.parent.parent
-urdf_filename = "talos_reduced.urdf"
-urdf_model_path = model_path / "talos_data/robots" / urdf_filename
+urdf_filename = "panda.urdf"
+urdf_model_path = model_path / "panda_description/urdf" / urdf_filename
 
 model, collision_model, visual_model = pin.buildModelsFromUrdf(
     urdf_model_path, mesh_dir, pin.JointModelFreeFlyer()

examples/viser-viewer.py

@@ -17,8 +17,8 @@
 mesh_dir = pinocchio_model_dir
 # urdf_filename = "talos_reduced.urdf"
 # urdf_model_path = join(join(model_path,"talos_data/robots"),urdf_filename)
-urdf_filename = "solo.urdf"
-urdf_model_path = model_path / "solo_description/robots" / urdf_filename
+urdf_filename = "tiago.urdf"
+urdf_model_path = model_path / "tiago_description/robots" / urdf_filename
 
 model, collision_model, visual_model = pin.buildModelsFromUrdf(
     urdf_model_path, mesh_dir, pin.JointModelFreeFlyer()