Save OBJ files for URDF export (#4)

* save obj files, run end-to-end test

* don't run e2e test by default

kol/mesh.py

@@ -0,0 +1,55 @@
+"""Defines utility functions for operations on meshes."""
+
+from pathlib import Path
+from typing import Literal, cast
+
+import stl.mesh
+
+MeshExt = Literal["stl", "obj"]
+
+
+def stl_to_obj(stl_path: str | Path, obj_path: str | Path) -> None:
+    """Converts an STL file to an OBJ file.
+
+    Args:
+        stl_path: The path to the STL file.
+        obj_path: The path to the OBJ file.
+    """
+    mesh = stl.mesh.Mesh.from_file(stl_path)
+
+    vertices: dict[tuple[float, float, float], int] = {}
+    faces: list[tuple[int, int, int]] = []
+    index = 1
+
+    # Process each triangle in the mesh
+    for i in range(len(mesh.vectors)):
+        face = []
+        for point in mesh.vectors[i]:
+            vertex = cast(tuple[float, float, float], tuple(point))
+            if vertex not in vertices:
+                vertices[vertex] = index
+                index += 1
+            face.append(vertices[vertex])
+        face_tuple = cast(tuple[int, int, int], tuple(face))
+        faces.append(face_tuple)
+
+    with open(obj_path, "w") as f:
+        for vertex, _ in sorted(vertices.items(), key=lambda x: x[1]):
+            f.write(f"v {' '.join(map(str, vertex))}\n")
+        for face_tuple in faces:
+            f.write(f"f {' '.join(map(str, face_tuple))}\n")
+
+
+def stl_to_fmt(stl_path: str | Path, output_path: str | Path) -> None:
+    stl_path = Path(stl_path)
+    ext = Path(output_path).suffix[1:]
+
+    match ext:
+        case "stl":
+            return
+
+        case "obj":
+            stl_to_obj(stl_path, output_path)
+
+        case _:
+            raise ValueError(f"Unsupported mesh format: {ext}")

kol/onshape/converter.py

@@ -19,6 +19,7 @@
 
 from kol import urdf
 from kol.geometry import apply_matrix_, inv_tf, rotation_matrix_to_euler_angles
+from kol.mesh import MeshExt, stl_to_fmt
 from kol.onshape.api import OnshapeApi
 from kol.onshape.client import OnshapeClient
 from kol.onshape.schema.assembly import (
@@ -97,6 +98,7 @@ class Converter:
             This is used to override the default joint velocity limits by
             matching the suffix of the joint name.
         disable_mimics: Whether to disable mimic joints.
+        mesh_ext: The extension of the mesh files to download.
     """
 
     def __init__(
@@ -109,6 +111,7 @@ def __init__(
         suffix_to_joint_effort: list[tuple[str, float]] = [],
         suffix_to_joint_velocity: list[tuple[str, float]] = [],
         disable_mimics: bool = False,
+        mesh_ext: MeshExt = "stl",
     ) -> None:
         super().__init__()
 
@@ -130,6 +133,7 @@ def __init__(
         self.suffix_to_joint_effort = [(k.lower().strip(), v) for k, v in suffix_to_joint_effort]
         self.suffix_to_joint_velocity = [(k.lower().strip(), v) for k, v in suffix_to_joint_velocity]
         self.disable_mimics = disable_mimics
+        self.mesh_ext = mesh_ext
 
         # Map containing all cached items.
         self.cache_map: dict[str, Any] = {}
@@ -595,18 +599,27 @@ def get_urdf_part(self, key: Key, joint: Joint | None = None) -> urdf.Link:
             stl_origin_to_part_tf = inv_tf(joint.child_entity.matedCS.part_to_mate_tf)
         self.stl_origin_to_part_tf[key] = stl_origin_to_part_tf
 
-        part_file_name = f"{part_name}{configuration_str}.stl"
+        part_file_name = f"{part_name}{configuration_str}.{self.mesh_ext}"
         part_file_path = self.mesh_dir / part_file_name
+
         if part_file_path.exists():
-            logger.info("Using cached STL file %s", part_file_path)
+            logger.info("Using cached file %s", part_file_path)
+
         else:
-            logger.info("Downloading STL file %s", part_file_path)
-            buffer = io.BytesIO()
-            self.api.download_stl(part, buffer)
-            buffer.seek(0)
-            mesh = stl.mesh.Mesh.from_file(None, fh=buffer)
-            mesh = apply_matrix_(mesh, stl_origin_to_part_tf)
-            mesh.save(part_file_path)
+            # Downloads the STL file.
+            part_file_path_stl = part_file_path.with_suffix(".stl")
+            if not part_file_path_stl.exists():
+                logger.info("Downloading file %s", part_file_path_stl)
+                buffer = io.BytesIO()
+                self.api.download_stl(part, buffer)
+                buffer.seek(0)
+                mesh = stl.mesh.Mesh.from_file(None, fh=buffer)
+                mesh = apply_matrix_(mesh, stl_origin_to_part_tf)
+                mesh.save(part_file_path_stl)
+
+            # Converts the mesh to the desired format.
+            logger.info("Converting STL file to %s", part_file_path)
+            stl_to_fmt(part_file_path_stl, part_file_path)
 
         # Move the mesh origin and dynamics from the part frame to the parent
         # joint frame (since URDF expects this by convention).

kol/scripts/get_urdf.py

@@ -3,12 +3,13 @@
 
 import argparse
 import logging
-from typing import Sequence
+from typing import Sequence, get_args
 
 import numpy as np
 
 from kol import urdf
 from kol.logging import configure_logging
+from kol.mesh import MeshExt
 from kol.onshape.converter import Converter
 
 
@@ -26,6 +27,7 @@ def main(args: Sequence[str] | None = None) -> None:
     parser.add_argument("--suffix-to-joint-effort", type=str, nargs="+", help="The suffix to joint effort mapping")
     parser.add_argument("--suffix-to-joint-velocity", type=str, nargs="+", help="The suffix to joint velocity mapping")
     parser.add_argument("--disable-mimics", action="store_true", help="Disable the mimic joints")
+    parser.add_argument("--mesh-ext", type=str, default="stl", choices=get_args(MeshExt), help="The mesh file format")
     parsed_args = parser.parse_args(args)
 
     configure_logging(level=logging.DEBUG if parsed_args.debug else logging.INFO)
@@ -60,6 +62,7 @@ def main(args: Sequence[str] | None = None) -> None:
         suffix_to_joint_effort=suffix_to_joint_effort,
         suffix_to_joint_velocity=suffix_to_joint_velocity,
         disable_mimics=parsed_args.disable_mimics,
+        mesh_ext=parsed_args.mesh_ext,
     ).save_urdf()
 
 

tests/test_e2e.py

@@ -0,0 +1,51 @@
+"""Runs an end-to-end test of the URDF exporter on the Stompy model."""
+
+from pathlib import Path
+from typing import get_args
+
+import pytest
+
+from kol import urdf
+from kol.logging import configure_logging
+from kol.mesh import MeshExt
+from kol.onshape.converter import Converter
+
+STOMPY_ONSHAPE_URL = (
+    "https://cad.onshape.com/documents/71f793a23ab7562fb9dec82d/w/"
+    "6160a4f44eb6113d3fa116cd/e/1a95e260677a2d2d5a3b1eb3"
+)
+
+
+@pytest.mark.skip(reason="This test is slow and requires an internet connection")
+def test_e2e(tmpdir: Path) -> None:
+    """Runs an end-to-end test of the URDF exporter on the Stompy model.
+
+    Args:
+        tmpdir: The temporary directory to save the URDF file.
+        ext: The mesh file format.
+    """
+    for mesh_ext in get_args(MeshExt):
+        Converter(
+            document_url=STOMPY_ONSHAPE_URL,
+            output_dir=tmpdir,
+            default_prismatic_joint_limits=urdf.JointLimits(10, 10, -10, 10),
+            default_revolute_joint_limits=urdf.JointLimits(10, 10, -10, 10),
+            suffix_to_joint_effort=[
+                ("dof_x4_h", 1.5),
+                ("dof_x4", 1.5),
+                ("dof_x6", 3),
+                ("dof_x8", 6),
+                ("dof_x10", 12),
+                ("knee_revolute", 13.9),
+                ("ankle_revolute", 6),
+            ],
+            suffix_to_joint_velocity=[],
+            disable_mimics=True,
+            mesh_ext=mesh_ext,
+        ).save_urdf()
+
+
+if __name__ == "__main__":
+    # python -m tests.test_e2e
+    configure_logging()
+    test_e2e(Path.cwd() / "test_e2e")

tests/test_mesh.py

@@ -0,0 +1,11 @@
+"""Tests mesh conversion options."""
+
+from pathlib import Path
+
+from kol.mesh import stl_to_fmt
+
+
+def test_stl_to_obj(tmpdir: Path) -> None:
+    obj_path = Path(tmpdir / "random.obj")
+    stl_to_fmt("tests/data/random.stl", obj_path)
+    assert obj_path.exists()