Merge pull request #33 from lorycontixd/add_mesh_support

Add support for meshes to primitive builders

.github/workflows/ci_cd.yml

@@ -122,7 +122,8 @@ jobs:
             pptree \
             idyntree \
             pytest \
-            robot_descriptions
+            robot_descriptions \
+            trimesh
             # pytest-icdiff \  # creates problems on macOS
           mamba install -y gz-sim7 idyntree
 

setup.cfg

@@ -61,6 +61,7 @@ install_requires =
     packaging
     resolve-robotics-uri-py
     scipy
+    trimesh
     xmltodict
 
 [options.extras_require]

src/rod/builder/primitives.py

@@ -1,4 +1,9 @@
 import dataclasses
+import pathlib
+from typing import Union
+
+import trimesh
+from numpy.typing import NDArray
 
 import rod
 from rod.builder.primitive_builder import PrimitiveBuilder
@@ -54,3 +59,52 @@ def _geometry(self) -> rod.Geometry:
         return rod.Geometry(
             cylinder=rod.Cylinder(radius=self.radius, length=self.length)
         )
+
+
+@dataclasses.dataclass
+class MeshBuilder(PrimitiveBuilder):
+    mesh_path: Union[str, pathlib.Path]
+    scale: NDArray
+
+    def __post_init__(self) -> None:
+        """
+        Post-initialization method for the class.
+        Loads the mesh from the specified file path and performs necessary checks.
+
+        Raises:
+            AssertionError: If the scale is not a 3D vector.
+            TypeError: If the mesh_path is not a str or pathlib.Path.
+        """
+
+        if isinstance(self.mesh_path, str):
+            extension = pathlib.Path(self.mesh_path).suffix
+        elif isinstance(self.mesh_path, pathlib.Path):
+            extension = self.mesh_path.suffix
+        else:
+            raise TypeError(
+                f"Expected str or pathlib.Path for mesh_path, got {type(self.mesh_path)}"
+            )
+
+        self.mesh: trimesh.base.Trimesh = trimesh.load(
+            str(self.mesh_path),
+            force="mesh",
+            file_type=extension,
+        )
+
+        assert self.scale.shape == (
+            3,
+        ), f"Scale must be a 3D vector, got {self.scale.shape}"
+
+    def _inertia(self) -> rod.Inertia:
+        inertia = self.mesh.moment_inertia
+        return rod.Inertia(
+            ixx=inertia[0, 0],
+            ixy=inertia[0, 1],
+            ixz=inertia[0, 2],
+            iyy=inertia[1, 1],
+            iyz=inertia[1, 2],
+            izz=inertia[2, 2],
+        )
+
+    def _geometry(self) -> rod.Geometry:
+        return rod.Geometry(mesh=rod.Mesh(uri=str(self.mesh_path), scale=self.scale))

tests/test_meshbuilder.py

@@ -0,0 +1,61 @@
+import os
+import pathlib
+import tempfile
+
+import numpy as np
+import trimesh
+
+from rod.builder.primitives import MeshBuilder
+
+
+def test_builder_creation():
+    mesh = trimesh.creation.box([1, 1, 1])
+
+    # Temporary write to file because rod Mesh works with uri
+    with tempfile.NamedTemporaryFile(suffix=".stl") as fp:
+        mesh.export(fp.name, file_type="stl")
+
+        builder = MeshBuilder(
+            name="test_mesh",
+            mesh_path=fp.name,
+            mass=1.0,
+            scale=np.array([1.0, 1.0, 1.0]),
+        )
+    assert (
+        builder.mesh.vertices.shape == mesh.vertices.shape
+    ), f"{builder.mesh.vertices.shape} != {mesh.vertices.shape}"
+    assert (
+        builder.mesh.faces.shape == mesh.faces.shape
+    ), f"{builder.mesh.faces.shape} != {mesh.faces.shape}"
+    assert (
+        builder.mesh.moment_inertia.all() == mesh.moment_inertia.all()
+    ), f"{builder.mesh.moment_inertia} != {mesh.moment_inertia}"
+    assert builder.mesh.volume == mesh.volume, f"{builder.mesh.volume} != {mesh.volume}"
+
+
+def test_builder_creation_custom_mesh():
+    # Create a custom mesh
+    vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]])
+    faces = np.array([[0, 1, 2], [0, 2, 3]])
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+
+    # Temporary write to file because rod Mesh works with uri
+    with tempfile.NamedTemporaryFile(suffix=".stl") as fp:
+        mesh.export(fp.name, file_type="stl")
+
+        builder = MeshBuilder(
+            name="test_mesh",
+            mesh_path=fp.name,
+            mass=1.0,
+            scale=np.array([1.0, 1.0, 1.0]),
+        )
+    assert (
+        builder.mesh.vertices.shape == mesh.vertices.shape
+    ), f"{builder.mesh.vertices.shape} != {mesh.vertices.shape}"
+    assert (
+        builder.mesh.faces.shape == mesh.faces.shape
+    ), f"{builder.mesh.faces.shape} != {mesh.faces.shape}"
+    assert (
+        builder.mesh.moment_inertia.all() == mesh.moment_inertia.all()
+    ), f"{builder.mesh.moment_inertia} != {mesh.moment_inertia}"
+    assert builder.mesh.volume == mesh.volume, f"{builder.mesh.volume} != {mesh.volume}"