Merge pull request #4 from readyplayerme/feat/group-seam-vertices

feat: identify overlapping vertices at UV seams

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

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):

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)