diff --git a/src/readyplayerme/meshops/mesh.py b/src/readyplayerme/meshops/mesh.py index 787ca28..2ef70cc 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 +from readyplayerme.meshops.types import Edges, IndexGroups, Indices, Mesh, Vertices def read_mesh(filename: str | Path) -> Mesh: @@ -47,3 +48,47 @@ def get_boundary_vertices(edges: Edges) -> Indices: unique_edges, edge_triangle_count = np.unique(sorted_edges, return_counts=True, axis=0) border_edge_indices = np.where(edge_triangle_count == 1)[0] return np.unique(unique_edges[border_edge_indices]) + + +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. + + :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. + """ + # Not using try / except because when using an index of -1 gets the last element and creates a false positive + if indices is None: + selected_vertices = vertices_pos + else: + if len(indices) == 0: + return [] + if np.any(indices < 0): + msg = "Negative index value is not allowed." + raise IndexError(msg) + + if np.max(indices) >= len(vertices_pos): + msg = "Index is out of bounds." + raise IndexError(msg) + + selected_vertices = 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 aee2060..d4189f7 100644 --- a/src/readyplayerme/meshops/types.py +++ b/src/readyplayerme/meshops/types.py @@ -5,10 +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) +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 eaaa655..1d09b17 100644 --- a/tests/readyplayerme/meshops/unit/test_mesh.py +++ b/tests/readyplayerme/meshops/unit/test_mesh.py @@ -8,7 +8,7 @@ import pytest from readyplayerme.meshops import mesh -from readyplayerme.meshops.types import Mesh +from readyplayerme.meshops.types import IndexGroups, Indices, Mesh, Vertices class TestReadMesh: @@ -57,3 +57,65 @@ def test_get_boundary_vertices(mock_mesh: Mesh): assert np.array_equiv( np.sort(boundary_vertices), [0, 2, 4, 6, 7, 9, 10] ), "The vertices returned by get_border_vertices do not match the expected vertices." + + +@pytest.mark.parametrize( + "vertices, indices, precision, expected", + [ + # Vertices from our mock mesh. + ("mock_mesh", np.array([0, 2, 4, 6, 7, 9, 10]), 0.1, [np.array([9, 10])]), + # Close positions, but with imprecision. + ( + np.array( + [ + [1.0, 1.0, 1.0], + [0.99998, 0.99998, 0.99998], + [0.49998, 0.5, 0.5], + [0.5, 0.5, 0.5], + [0.50001, 0.50001, 0.50001], + ] + ), + 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), 0.1, []), + ], +) +def test_get_overlapping_vertices( + vertices: Vertices, indices: Indices, precision: float, expected: IndexGroups, 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, precision) + + assert len(grouped_vertices) == len(expected), "Number of groups doesn't match expected" + 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}" + + +@pytest.mark.parametrize( + "indices", + [ + # Case with index out of bounds (higher than max) + np.array([0, 3], dtype=np.uint16), + # Case with index out of bounds (negative index) + np.array([0, -1], dtype=np.int32), # Using int32 to allow negative values + ], +) +def test_get_overlapping_vertices_error_handling(indices): + """Test that get_overlapping_vertices function raises an exception for out of bounds indices.""" + vertices = np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]) + with pytest.raises(IndexError): + mesh.get_overlapping_vertices(vertices, indices)