diff --git a/src/readyplayerme/meshops/mesh.py b/src/readyplayerme/meshops/mesh.py index 315b630..6c3caf7 100644 --- a/src/readyplayerme/meshops/mesh.py +++ b/src/readyplayerme/meshops/mesh.py @@ -4,8 +4,9 @@ import numpy as np import trimesh +from scipy.spatial import cKDTree -from readyplayerme.meshops.types import Edges, Indices, Mesh, VariableLengthArrays, Vertices +from readyplayerme.meshops.types import Edges, IndexGroups, Indices, Mesh, Vertices def read_mesh(filename: str | Path) -> Mesh: @@ -49,27 +50,31 @@ def get_boundary_vertices(edges: Edges) -> Indices: return np.unique(unique_edges[border_edge_indices]) -def get_overlapping_vertices(vertices_pos: Vertices, indices: Indices | None = None) -> VariableLengthArrays: +def get_overlapping_vertices( + vertices_pos: Vertices, indices: Indices | None = None, tolerance: float = 0.00001 +) -> IndexGroups: """Return the indices of the vertices grouped by the same position. - Vertices that have the same position belong to a seam. - - - :param vertices: All the vertices of the mesh. - :param indices: Vertex indices. - :return: A list of grouped border vertices that share position. + :param vertices_pos: All the vertices of the mesh. + :param indices: Vertex indices. + :param precision: Tolerance for considering positions as overlapping. + :return: A list of grouped vertices that share position. """ - if indices is None: - indices = np.arange(len(vertices_pos)) - - vertex_positions = vertices_pos[indices] - rounded_positions = np.round(vertex_positions, decimals=5) - structured_positions = np.core.records.fromarrays( - rounded_positions.transpose(), names="x, y, z", formats="f8, f8, f8" - ) - unique_positions, local_indices = np.unique(structured_positions, return_inverse=True) - grouped_indices = [ - indices[local_indices == i] for i in range(len(unique_positions)) if (local_indices == i).sum() > 1 - ] + selected_vertices = vertices_pos if indices is None else vertices_pos[indices] + + tree = cKDTree(selected_vertices) + + grouped_indices = [] + processed = set() + for idx, vertex in enumerate(selected_vertices): + if idx not in processed: + # Find all points within the tolerance distance + neighbors = tree.query_ball_point(vertex, tolerance) + if len(neighbors) > 1: # Include only groups with multiple vertices + # Translate to original indices if needed + group = np.array(neighbors, dtype=np.uint32) if indices is None else indices[neighbors] + grouped_indices.append(group) + # Mark these points as processed + processed.update(neighbors) return grouped_indices diff --git a/src/readyplayerme/meshops/types.py b/src/readyplayerme/meshops/types.py index dc8535e..d4189f7 100644 --- a/src/readyplayerme/meshops/types.py +++ b/src/readyplayerme/meshops/types.py @@ -5,11 +5,11 @@ import numpy.typing as npt # trimesh uses int64 and float64 for its arrays. -Indices: TypeAlias = npt.NDArray[np.int32] | npt.NDArray[np.int64] # Shape (i,) +Indices: TypeAlias = npt.NDArray[np.uint32] | npt.NDArray[np.uint64] # Shape (i,) Vertices: TypeAlias = npt.NDArray[np.float32] | npt.NDArray[np.float64] # Shape (v, 3) Edges: TypeAlias = npt.NDArray[np.int32] | npt.NDArray[np.int64] # Shape (e, 2) Faces: TypeAlias = npt.NDArray[np.int32] | npt.NDArray[np.int64] # Shape (f, 3) -VariableLengthArrays: TypeAlias = list[npt.NDArray[np.int64]] +IndexGroups: TypeAlias = list[npt.NDArray[np.uint32]] class Mesh(Protocol): diff --git a/tests/readyplayerme/meshops/unit/test_mesh.py b/tests/readyplayerme/meshops/unit/test_mesh.py index 0a5b39c..7ea892e 100644 --- a/tests/readyplayerme/meshops/unit/test_mesh.py +++ b/tests/readyplayerme/meshops/unit/test_mesh.py @@ -60,10 +60,10 @@ def test_get_boundary_vertices(mock_mesh: Mesh): @pytest.mark.parametrize( - "vertices, indices, expected", + "vertices, indices, precision, expected", [ # Vertices from our mock mesh. - ("mock_mesh", np.array([0, 2, 4, 6, 7, 9, 10]), [(np.array([9, 10]))]), + ("mock_mesh", np.array([0, 2, 4, 6, 7, 9, 10]), 0.1, [np.array([9, 10])]), # Close positions, but with imprecision. ( np.array( @@ -76,27 +76,30 @@ def test_get_boundary_vertices(mock_mesh: Mesh): ] ), np.array([0, 1, 2, 3, 4]), + 0.0001, [np.array([0, 1]), np.array([2, 3, 4])], ), # Overlapping vertices, None indices given. ( np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), None, + 0.1, [np.array([0, 1])], ), # Overlapping vertices, but empty indices given. - (np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), np.array([], dtype=np.int32), []), + (np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), np.array([], dtype=np.int32), 0.1, []), ], ) -def test_overlapping_vertices(vertices: Vertices, indices: Indices, expected: Indices, request: pytest.FixtureRequest): +def test_get_overlapping_vertices( + vertices: Vertices, indices: Indices, precision: float, expected: Indices, request: pytest.FixtureRequest +): """Test the get_overlapping_vertices functions returns the expected indices groups.""" # Get vertices from the fixture if one is given. if isinstance(vertices, str) and vertices == "mock_mesh": vertices = request.getfixturevalue("mock_mesh").vertices - grouped_vertices = mesh.get_overlapping_vertices(vertices, indices) + grouped_vertices = mesh.get_overlapping_vertices(vertices, indices, precision) assert len(grouped_vertices) == len(expected), "Number of groups doesn't match expected" - assert np.array_equiv( - grouped_vertices, expected - ), "The vertices returned by get_overlapping_vertices do not match the expected vertices." + for group, exp_group in zip(grouped_vertices, expected, strict=False): + assert np.array_equal(group, exp_group), f"Grouped vertices {group} do not match expected {exp_group}"