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,16 @@
 
 import numpy as np
 
+from pathlib import Path
+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,22 +219,111 @@ 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.
+        Load a mesh from a file
         """
+        print("on passe la")
+        mesh_path = Path(mesh_path)
+        ext = mesh_path.suffix.lower()
+
+        if ext == ".dae":
+            try:
+                mesh_collada = collada.Collada(mesh_path)
+
+                # If there are more geometries than effects, fall back to trimesh
+                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:
+                        # Extract vertices and indices from Collada geometry
+                        vertices = geometry.primitives[0].sources["VERTEX"][0][4].data
+                        indices = geometry.primitives[0].indices
+
+                        # Handle different shapes of indices array
+                        if indices.ndim == 3:
+                            faces = indices[:, :, 0]
+                        elif indices.ndim in (1, 2):
+                            faces = indices.reshape(-1, 3)
+                        else:
+                            faces = indices.reshape(-1, 3)
+
+                        # If after reshaping the last dimension isn't 3, try reshaping again
+                        if faces.shape[-1] != 3 and faces.size % 3 == 0:
+                            faces = faces.reshape(-1, 3)
+
+                        mesh_color = getattr(effect, "diffuse", None)
+
+                        # Add mesh with appropriate color
+                        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:
+                            # No color available, fallback to trimesh loader
+                            mesh = trimesh.load_mesh(mesh_path)
+                            frame = self.viewer.scene.add_mesh_trimesh(f"{name}_{i}", mesh)
+
+                        # If for some reason the mesh wasn't added, use a tiny placeholder box
+                        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:
+                        # Failed to process this geometry manually, fallback to trimesh
+                        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:
+                            # Totally failed, skip this geometry
+                            pass
+
+                return frames[0] if frames else None
+
+            except Exception:
+                # Failed to load as Collada, fallback to trimesh for other formats or malformed .dae
+                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]
+                    )
+
+        # For non-Collada files
         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):
         """
         Load a mesh from triangles stored inside a hppfcl.Convex.

examples/viser-viewer.py

@@ -20,6 +20,7 @@
 urdf_filename = "solo.urdf"
 urdf_model_path = model_path / "solo_description/robots" / urdf_filename
 
+
 model, collision_model, visual_model = pin.buildModelsFromUrdf(
     urdf_model_path, mesh_dir, pin.JointModelFreeFlyer()
 )