From 3edae89ae363bf254a8489a0c50e8395e600bb59 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Fri, 15 Nov 2024 17:44:12 +0100 Subject: [PATCH] Merge pull request #156 from lorycontixd/add_mesh_support commit de410bf6102772f7078abac0c5ae77333362a0d3 Author: Lorenzo Conti Date: Thu Jul 18 12:17:22 2024 +0200 Added string magics on wrapping methods - Added string magics on mesh wrapping methods for easier logging - Added log to indicate number of extracted points - Added extra check on objectmapping method commit 1533f7016a3c562961b625bbf86d3550201d83e8 Author: Lorenzo Conti Date: Fri Jul 5 14:53:48 2024 +0200 Moved from jax.numpy to numpy in PlaneTerrain __eq__ magic to bypass TracerError commit e3f167a55a4a28b3d690b8d42c3bca9908d8280d Merge: 8fa9adc fe2616c Author: Filippo Luca Ferretti <102977828+flferretti@users.noreply.github.com> Date: Fri Nov 15 17:35:46 2024 +0100 Merge pull request #156 from lorycontixd/add_mesh_support Add mesh support commit fe2616cbcf7c66d6c86b43988c9fe061d264886d Author: Lorenzo Conti Date: Fri Nov 15 17:18:07 2024 +0100 Removed whitespaces commit 2d46a704d6784bcce068daf71b5d20d93e558b67 Author: Lorenzo Conti Date: Fri Nov 15 17:16:08 2024 +0100 Fixed minor commenting format commit eff915d1c98d9d5692188919d67c80ce1764b154 Author: Lorenzo Conti Date: Fri Nov 15 17:14:35 2024 +0100 Removed unused function commit 1738a879cf9ea4f4283431ea38a0cda510d85420 Author: Lorenzo Conti Date: Fri Nov 15 17:12:19 2024 +0100 Removed unused dependency from pyprojecj commit bb21a1619289de41570fd825d5dfcbc2b102da48 Author: Lorenzo Conti Date: Fri Nov 15 17:10:09 2024 +0100 Precommit fix commit 495b799766fb760b6d71d2511b0632da3845415a Author: Lorenzo Conti <56968043+lorycontixd@users.noreply.github.com> Date: Fri Nov 15 16:43:42 2024 +0100 Apply suggestions from code review Co-authored-by: Filippo Luca Ferretti <102977828+flferretti@users.noreply.github.com> commit 90451e8ec1307ba65ac8b9e732328c5552948baa Author: Lorenzo Conti Date: Fri Nov 15 16:41:36 2024 +0100 Removed extra search paths in ergocub model building commit e76b0b9fd8ab7e4ef9bd10d970da122ee9dfcdcd Author: Lorenzo Conti Date: Fri Nov 15 16:16:55 2024 +0100 Added int casting on mesh_enabled flag commit 8a2934971b3e4b9c3924f904846e6c4c5c9d0bf9 Author: Lorenzo Conti <56968043+lorycontixd@users.noreply.github.com> Date: Fri Nov 15 15:27:02 2024 +0100 Updated variable names Co-authored-by: Filippo Luca Ferretti <102977828+flferretti@users.noreply.github.com> commit 8d0380e4db2532f0d1f8e2f4f516652d282159c1 Author: Lorenzo Conti Date: Fri Nov 15 16:14:26 2024 +0100 Added experimental feature warning for mesh parsing commit ad30c29b9097edcc2afebe7d71b0f4dd1bd14318 Author: Lorenzo Conti Date: Fri Nov 15 15:23:55 2024 +0100 Addressed reviews - Removed unused dependencies in pyproject.toml - Moved from warning to exception if passed mesh is empty Co-authored-by: Filippo Luca Ferretti commit 9238af5c118cb4075850c68434fbcba20a03b763 Author: Lorenzo Conti Date: Fri Nov 15 12:10:38 2024 +0100 Fixed error on array sorting and relative test commit fec016a7049ac81fd0bc710265361f83807f2755 Author: Lorenzo Conti Date: Fri Nov 15 11:41:24 2024 +0100 Implemented reviews commit d2a5e89e9d2ec82b4d9ccb559570539c0f93f822 Author: Lorenzo Conti Date: Thu Jul 18 12:17:22 2024 +0200 Added string magics on wrapping methods - Added string magics on mesh wrapping methods for easier logging - Added log to indicate number of extracted points - Added extra check on objectmapping method commit 98012e59a471ce3063bc9184ad6d2dfaea7d667c Author: Lorenzo Conti Date: Thu Nov 14 12:02:43 2024 +0100 Added docstrings to mesh wrapping algorithms commit 148fcfee6875b72fb20f8e6a4a5e097a61bcf66d Author: Lorenzo Conti Date: Thu Jul 18 12:17:22 2024 +0200 Added string magics on wrapping methods - Added string magics on mesh wrapping methods for easier logging - Added log to indicate number of extracted points - Added extra check on objectmapping method commit 4f6cf0a8fbe906bbf6092630a317062ac85904e0 Author: Lorenzo Conti Date: Thu Jul 18 12:11:56 2024 +0200 Removed wrong point selection & added logs - Removed a line that would always set the extracted points to the vertices - Added a few debug lines using logger commit f8ecf24a6cebf67105bde8e914fa8762fc66b1b1 Author: Lorenzo Conti Date: Thu Jul 18 12:01:34 2024 +0200 Removed leftover parameters on create_mesh_collision commit 6acb19fa7942ab24a6a4fc6422a635b6490e26cf Author: Lorenzo Conti Date: Thu Jul 18 11:55:01 2024 +0200 Run pre-commit commit 72ce440705a0c9d9efa973ef2355ac4df1b61ab1 Author: Lorenzo Conti Date: Thu Jul 18 11:50:37 2024 +0200 Restructured mesh mapping methods to follow inheritance - Redefined methods using classes - Adapted rod parser to new structure - Reimplemented tests with new structure commit 2d07347fa6d0690d864666ac53e6d109868adcda Author: Lorenzo Conti Date: Thu Jul 18 10:45:09 2024 +0200 Renamed some parameters - Renamed some function paramaeters - Added a few tests TODO: migrate MeshMapping static class to inheritance structure commit 8058b7c9c37061252fae838bd8d12105815479f7 Author: Lorenzo Conti Date: Thu Jul 18 10:25:51 2024 +0200 New mesh wrapping algorithms with relative tests - New mesh wrapping algorithms (mesh decimation, object mapping, aap, select points over axis) - Implemented tests of above except first algorithm - Updated manifold3d dependency (used in object mapping) - Restructured meshes module commit 22da2cde555816f30a26da7f5575e4ee25715ee0 Author: Lorenzo Conti Date: Wed Jul 17 23:00:18 2024 +0200 New mesh wrapping algorithms - Implemented AAP algorithm - Restructured collision parsing to accept the new algorithms - Wrote tests for AAP algorithm - Updated JaxSim dependencies commit d434e445032e0322f0018bf0c64afd3fb207c9c9 Author: Lorenzo Conti Date: Wed Jul 17 17:56:46 2024 +0200 Implemented structure for new mesh wrapping algorithms commit f9475b0e4c2a01213a73673f0b924b745d7387f1 Author: Lorenzo Conti Date: Wed Jul 17 17:55:25 2024 +0200 First draft of new mesh wrapping algorithms commit 5af51da739b4ca3b3c792b0d9b01406e5ebfdcb7 Author: Lorenzo Conti Date: Fri Jul 5 16:08:34 2024 +0200 Implemented initial reviews from https://github.com/ami-iit/jaxsim/pull/156 - Restructured mesh_collision creation method - Removed unnecessary hash inheritance commit ed51e85e0a650cfb179a59b0f6761f5503e3da13 Author: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri Jul 5 12:55:11 2024 +0000 [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci commit cf003352385c4cb4905b0f9b24311664aca5628b Author: Lorenzo Conti Date: Thu Jun 20 10:39:42 2024 +0200 Removed unused import in conftest commit 94c2f81b77794177fa88254f44fba7540beb7f0f Author: Lorenzo Conti Date: Thu Jun 20 10:37:23 2024 +0200 Fixed typo on logging message commit d7fd1b3f33eddb2d4533268528b71e6a53890d2b Author: Lorenzo Conti Date: Thu Jun 20 10:34:52 2024 +0200 Removed unused lines in conftest commit 07b84025b93c3a41e6087ce97be358c3b7c9e9a6 Author: Filippo Luca Ferretti Date: Thu Jun 13 12:34:26 2024 +0200 Update `__eq__` magic and type hints commit bcf5e48cb0cc09d80681d376fc00fc7b9093ace2 Author: Lorenzo Conti Date: Tue May 21 14:59:23 2024 +0200 Pre-commit commit b0ddec10fdf46cba91d37972eb0e0cfb85cc2a75 Author: Lorenzo Conti Date: Tue May 21 14:57:30 2024 +0200 Address reviews - Remove leftover comments - MeshMappingMethods inherit from IntEnum instead of Enum - Added center comparison for MeshCollision object commit 644ce435f4a012dd738bfbd1ee2c8600fa85e5ae Author: Lorenzo Conti Date: Tue May 21 14:51:08 2024 +0200 Implemented UniformSurfaceSampling for mesh point wrapping commit b15ea5514d6a8ae1603f527046d78763ae050122 Author: Lorenzo Conti Date: Tue May 21 10:11:55 2024 +0200 Moved mesh parsing logic inside mesh collision function commit 2f25c0ea3ef04113fa40192c5b1afcc0534873d3 Author: Lorenzo Conti Date: Tue May 21 10:10:49 2024 +0200 Added trimesh dependecy for conda-forge commit 46f7164cae8a269df9cf934511c1036872e191f3 Author: Lorenzo Conti Date: Mon May 20 18:11:36 2024 +0200 Address to reviews: - Moved mesh parsing logic inside method for creating mesh collisions - Removed vs code settings - Removed empty mesh parsing test commit 990ea06af346228bbb6bce6ca608c0412cd25463 Author: Lorenzo Conti Date: Thu May 16 14:34:48 2024 +0200 Added `networkx` as testing dependency commit 2c5f78fb769387a95c68e8869c322de498f73986 Author: Filippo Luca Ferretti Date: Mon May 20 16:16:41 2024 +0200 Skip loading empty meshes Co-authored-by: Lorenzo Conti commit 6e968fce7bd9b1ebff5643c3d51b8ed85e4e6a35 Author: Filippo Luca Ferretti Date: Fri May 17 10:22:23 2024 +0200 Use already existing env var to solve mesh URIs commit 1663a1f2e82ce9de2adea15492a13101f541cfb9 Author: Filippo Luca Ferretti Date: Tue May 14 15:35:49 2024 +0200 Format and lint commit 5d16532820332b466b363d939f0aac47dff929d0 Author: Lorenzo Conti Date: Mon May 13 19:42:47 2024 +0200 Initial version of mesh support commit d07e4d178e3225085b75db47ce001ccfb5445b93 Author: Filippo Luca Ferretti Date: Mon May 13 19:40:27 2024 +0200 Set env var when parsing from `robot-descriptions` commit 393663871dced8301d2b18615a04f70c777c051b Author: Lorenzo Conti Date: Fri Jul 5 14:53:48 2024 +0200 Moved from jax.numpy to numpy in PlaneTerrain __eq__ magic to bypass TracerError --- environment.yml | 1 + pyproject.toml | 3 +- src/jaxsim/parsers/descriptions/__init__.py | 8 +- src/jaxsim/parsers/descriptions/collision.py | 19 ++++ src/jaxsim/parsers/rod/meshes.py | 104 +++++++++++++++++++ src/jaxsim/parsers/rod/parser.py | 12 +++ src/jaxsim/parsers/rod/utils.py | 53 ++++++++++ tests/test_meshes.py | 100 ++++++++++++++++++ 8 files changed, 298 insertions(+), 2 deletions(-) create mode 100644 src/jaxsim/parsers/rod/meshes.py create mode 100644 tests/test_meshes.py diff --git a/environment.yml b/environment.yml index bc78942a2..2603b0c82 100644 --- a/environment.yml +++ b/environment.yml @@ -15,6 +15,7 @@ dependencies: - pptree - qpax - rod >= 0.3.3 + - trimesh - typing_extensions # python<3.12 # ==================================== # Optional dependencies from setup.cfg diff --git a/pyproject.toml b/pyproject.toml index f7ca141ec..1209baabe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dependencies = [ "qpax", "rod >= 0.3.3", "typing_extensions ; python_version < '3.12'", + "trimesh", ] [project.optional-dependencies] @@ -67,7 +68,7 @@ testing = [ "idyntree >= 12.2.1", "pytest >=6.0", "pytest-icdiff", - "robot-descriptions", + "robot-descriptions" ] viz = [ "lxml", diff --git a/src/jaxsim/parsers/descriptions/__init__.py b/src/jaxsim/parsers/descriptions/__init__.py index 9e180c155..ff3bf631d 100644 --- a/src/jaxsim/parsers/descriptions/__init__.py +++ b/src/jaxsim/parsers/descriptions/__init__.py @@ -1,4 +1,10 @@ -from .collision import BoxCollision, CollidablePoint, CollisionShape, SphereCollision +from .collision import ( + BoxCollision, + CollidablePoint, + CollisionShape, + MeshCollision, + SphereCollision, +) from .joint import JointDescription, JointGenericAxis, JointType from .link import LinkDescription from .model import ModelDescription diff --git a/src/jaxsim/parsers/descriptions/collision.py b/src/jaxsim/parsers/descriptions/collision.py index 31ae17b97..7815af488 100644 --- a/src/jaxsim/parsers/descriptions/collision.py +++ b/src/jaxsim/parsers/descriptions/collision.py @@ -154,3 +154,22 @@ def __eq__(self, other: BoxCollision) -> bool: return False return hash(self) == hash(other) + + +@dataclasses.dataclass +class MeshCollision(CollisionShape): + center: jtp.VectorLike + + def __hash__(self) -> int: + return hash( + ( + hash(tuple(self.center.tolist())), + hash(self.collidable_points), + ) + ) + + def __eq__(self, other: MeshCollision) -> bool: + if not isinstance(other, MeshCollision): + return False + + return hash(self) == hash(other) diff --git a/src/jaxsim/parsers/rod/meshes.py b/src/jaxsim/parsers/rod/meshes.py new file mode 100644 index 000000000..9d1ada7b0 --- /dev/null +++ b/src/jaxsim/parsers/rod/meshes.py @@ -0,0 +1,104 @@ +import numpy as np +import trimesh + +VALID_AXIS = {"x": 0, "y": 1, "z": 2} + + +def extract_points_vertices(mesh: trimesh.Trimesh) -> np.ndarray: + """ + Extracts the vertices of a mesh as points. + """ + return mesh.vertices + + +def extract_points_random_surface_sampling(mesh: trimesh.Trimesh, n) -> np.ndarray: + """ + Extracts N random points from the surface of a mesh. + + Args: + mesh: The mesh from which to extract points. + n: The number of points to extract. + + Returns: + The extracted points (N x 3 array). + """ + + return mesh.sample(n) + + +def extract_points_uniform_surface_sampling( + mesh: trimesh.Trimesh, n: int +) -> np.ndarray: + """ + Extracts N uniformly sampled points from the surface of a mesh. + + Args: + mesh: The mesh from which to extract points. + n: The number of points to extract. + + Returns: + The extracted points (N x 3 array). + """ + + return trimesh.sample.sample_surface_even(mesh=mesh, count=n)[0] + + +def extract_points_select_points_over_axis( + mesh: trimesh.Trimesh, axis: str, direction: str, n: int +) -> np.ndarray: + """ + Extracts N points from a mesh along a specified axis. The points are selected based on their position along the axis. + + Args: + mesh: The mesh from which to extract points. + axis: The axis along which to extract points. + direction: The direction along the axis from which to extract points. Valid values are "higher" and "lower". + n: The number of points to extract. + + Returns: + The extracted points (N x 3 array). + """ + + dirs = {"higher": np.s_[-n:], "lower": np.s_[:n]} + arr = mesh.vertices + + # Sort rows lexicographically first, then columnar. + arr.sort(axis=0) + sorted_arr = arr[dirs[direction]] + return sorted_arr + + +def extract_points_aap( + mesh: trimesh.Trimesh, + axis: str, + upper: float | None = None, + lower: float | None = None, +) -> np.ndarray: + """ + Extracts points from a mesh along a specified axis within a specified range. The points are selected based on their position along the axis. + + Args: + mesh: The mesh from which to extract points. + axis: The axis along which to extract points. + upper: The upper bound of the range. + lower: The lower bound of the range. + + Returns: + The extracted points (N x 3 array). + + Raises: + AssertionError: If the lower bound is greater than the upper bound. + """ + + # Check bounds. + upper = upper if upper is not None else np.inf + lower = lower if lower is not None else -np.inf + assert lower < upper, "Invalid bounds for axis-aligned plane" + + # Logic. + points = mesh.vertices[ + (mesh.vertices[:, VALID_AXIS[axis]] >= lower) + & (mesh.vertices[:, VALID_AXIS[axis]] <= upper) + ] + + return points diff --git a/src/jaxsim/parsers/rod/parser.py b/src/jaxsim/parsers/rod/parser.py index 8d359c6df..fc23420ae 100644 --- a/src/jaxsim/parsers/rod/parser.py +++ b/src/jaxsim/parsers/rod/parser.py @@ -334,6 +334,18 @@ def extract_model_data( collisions.append(sphere_collision) + if collision.geometry.mesh is not None and int( + os.environ.get("JAXSIM_COLLISION_MESH_ENABLED", "0") + ): + logging.warning("Mesh collision support is still experimental.") + mesh_collision = utils.create_mesh_collision( + collision=collision, + link_description=links_dict[link.name], + method=utils.meshes.extract_points_vertices, + ) + + collisions.append(mesh_collision) + return SDFData( model_name=sdf_model.name, link_descriptions=links, diff --git a/src/jaxsim/parsers/rod/utils.py b/src/jaxsim/parsers/rod/utils.py index dd83b1dde..5950a0fa8 100644 --- a/src/jaxsim/parsers/rod/utils.py +++ b/src/jaxsim/parsers/rod/utils.py @@ -1,12 +1,21 @@ import os +import pathlib +from collections.abc import Callable +from typing import TypeVar import numpy as np import numpy.typing as npt import rod +import trimesh +from rod.utils.resolve_uris import resolve_local_uri import jaxsim.typing as jtp +from jaxsim import logging from jaxsim.math import Adjoint, Inertia from jaxsim.parsers import descriptions +from jaxsim.parsers.rod import meshes + +MeshMappingMethod = TypeVar("MeshMappingMethod", bound=Callable[..., npt.NDArray]) def from_sdf_inertial(inertial: rod.Inertial) -> jtp.Matrix: @@ -202,3 +211,47 @@ def fibonacci_sphere(samples: int) -> npt.NDArray: return descriptions.SphereCollision( collidable_points=collidable_points, center=center_wrt_link ) + + +def create_mesh_collision( + collision: rod.Collision, + link_description: descriptions.LinkDescription, + method: MeshMappingMethod = None, +) -> descriptions.MeshCollision: + + file = pathlib.Path(resolve_local_uri(uri=collision.geometry.mesh.uri)) + _file_type = file.suffix.replace(".", "") + mesh = trimesh.load_mesh(file, file_type=_file_type) + + if mesh.is_empty: + raise RuntimeError(f"Failed to process '{file}' with trimesh") + + mesh.apply_scale(collision.geometry.mesh.scale) + logging.info( + msg=f"Loading mesh {collision.geometry.mesh.uri} with scale {collision.geometry.mesh.scale}, file type '{_file_type}'" + ) + + if method is None: + method = meshes.VertexExtraction() + logging.debug("Using default Vertex Extraction method for mesh wrapping") + else: + logging.debug(f"Using method {method} for mesh wrapping") + + points = method(mesh=mesh) + logging.debug(f"Extracted {len(points)} points from mesh") + + W_H_L = collision.pose.transform() if collision.pose is not None else np.eye(4) + + # Extract translation from transformation matrix + W_p_L = W_H_L[:3, 3] + mesh_points_wrt_link = points @ W_H_L[:3, :3].T + W_p_L + collidable_points = [ + descriptions.CollidablePoint( + parent_link=link_description, + position=point, + enabled=True, + ) + for point in mesh_points_wrt_link + ] + + return descriptions.MeshCollision(collidable_points=collidable_points, center=W_p_L) diff --git a/tests/test_meshes.py b/tests/test_meshes.py new file mode 100644 index 000000000..58fcb9827 --- /dev/null +++ b/tests/test_meshes.py @@ -0,0 +1,100 @@ +import trimesh + +from jaxsim.parsers.rod import meshes + + +def test_mesh_wrapping_vertex_extraction(): + """ + Test the vertex extraction method on different meshes. + 1. A simple box + 2. A sphere + """ + + # Test 1: A simple box. + # First, create a box with origin at (0,0,0) and extents (3,3,3), + # i.e. points span from -1.5 to 1.5 on the axis. + mesh = trimesh.creation.box( + extents=[3.0, 3.0, 3.0], + ) + points = meshes.extract_points_vertices(mesh=mesh) + assert len(points) == len(mesh.vertices) + + # Test 2: A sphere. + # The sphere is centered at the origin and has a radius of 1.0. + mesh = trimesh.creation.icosphere(subdivisions=4, radius=1.0) + points = meshes.extract_points_vertices(mesh=mesh) + assert len(points) == len(mesh.vertices) + + +def test_mesh_wrapping_aap(): + """ + Test the AAP wrapping method on different meshes. + 1. A simple box + 1.1: Remove all points above x=0.0 + 1.2: Remove all points below y=0.0 + 2. A sphere + """ + + # Test 1.1: Remove all points above x=0.0. + # The expected result is that the number of points is halved. + # First, create a box with origin at (0,0,0) and extents (3,3,3), + # i.e. points span from -1.5 to 1.5 on the axis. + mesh = trimesh.creation.box(extents=[3.0, 3.0, 3.0]) + points = meshes.extract_points_aap(mesh=mesh, axis="x", lower=0.0) + assert len(points) == len(mesh.vertices) // 2 + assert all(points[:, 0] > 0.0) + + # Test 1.2: Remove all points below y=0.0. + # The expected result is that the number of points is halved. + points = meshes.extract_points_aap(mesh=mesh, axis="y", upper=0.0) + assert len(points) == len(mesh.vertices) // 2 + assert all(points[:, 1] < 0.0) + + # Test 2: A sphere. + # The sphere is centered at the origin and has a radius of 1.0. + # Points are expected to be halved. + mesh = trimesh.creation.icosphere(subdivisions=4, radius=1.0) + + # Remove all points above y=0.0. + points = meshes.extract_points_aap(mesh=mesh, axis="y", lower=0.0) + assert all(points[:, 1] >= 0.0) + assert len(points) < len(mesh.vertices) + + +def test_mesh_wrapping_points_over_axis(): + """ + Test the points over axis method on different meshes. + 1. A simple box + 1.1: Select 10 points from the lower end of the x-axis + 1.2: Select 10 points from the higher end of the y-axis + 2. A sphere + """ + + # Test 1.1: Remove 10 points from the lower end of the x-axis. + # First, create a box with origin at (0,0,0) and extents (3,3,3), + # i.e. points span from -1.5 to 1.5 on the axis. + mesh = trimesh.creation.box(extents=[3.0, 3.0, 3.0]) + points = meshes.extract_points_select_points_over_axis( + mesh=mesh, axis="x", direction="lower", n=4 + ) + assert len(points) == 4 + assert all(points[:, 0] < 0.0) + + # Test 1.2: Select 10 points from the higher end of the y-axis. + points = meshes.extract_points_select_points_over_axis( + mesh=mesh, axis="y", direction="higher", n=4 + ) + assert len(points) == 4 + assert all(points[:, 1] > 0.0) + + # Test 2: A sphere. + # The sphere is centered at the origin and has a radius of 1.0. + mesh = trimesh.creation.icosphere(subdivisions=4, radius=1.0) + sphere_n_vertices = len(mesh.vertices) + + # Select 10 points from the higher end of the z-axis. + points = meshes.extract_points_select_points_over_axis( + mesh=mesh, axis="z", direction="higher", n=sphere_n_vertices // 2 + ) + assert len(points) == sphere_n_vertices // 2 + assert all(points[:, 2] >= 0.0)