From 050f650ae84b07fb430617a8db65d257e5df129a Mon Sep 17 00:00:00 2001 From: Krzysztof Chalupka Date: Mon, 11 Apr 2022 16:27:53 -0700 Subject: [PATCH] Submesh 3/n: Add submeshing functionality Summary: Copypasting the docstring: ``` Split a mesh into submeshes, defined by face indices of the original Meshes object. Args: face_indices: Let the original mesh have verts_list() of length N. Can be either - List of length N. The n-th element is a list of length num_submeshes_n (empty lists are allowed). Each element of the n-th sublist is a LongTensor of length num_faces. - List of length N. The n-th element is a possibly empty padded LongTensor of shape (num_submeshes_n, max_num_faces). Returns: Meshes object with selected submeshes. The submesh tensors are cloned. Currently submeshing only works with no textures or with the TexturesVertex texture. Example: Take a Meshes object `cubes` with 4 meshes, each a translated cube. Then: * len(cubes) is 4, len(cubes.verts_list()) is 4, len(cubes.faces_list()) is 4, * [cube_verts.size for cube_verts in cubes.verts_list()] is [8, 8, 8, 8], * [cube_faces.size for cube_faces in cubes.faces_list()] if [6, 6, 6, 6], Now let front_facet, top_and_bottom, all_facets be LongTensors of sizes (2), (4), and (12), each picking up a number of facets of a cube by specifying the appropriate triangular faces. Then let `subcubes = cubes.submeshes([[front_facet, top_and_bottom], [], [all_facets], []])`. * len(subcubes) is 3. * subcubes[0] is the front facet of the cube contained in cubes[0]. * subcubes[1] is a mesh containing the (disconnected) top and bottom facets of cubes[0]. * subcubes[2] is a clone of cubes[2]. * There are no submeshes of cubes[1] and cubes[3] in subcubes. * subcubes[0] and subcubes[1] are not watertight. subcubes[2] is. ``` Reviewed By: bottler Differential Revision: D35440657 fbshipit-source-id: 8a6d2d300ce226b5b9eb440688528b5e795195a1 --- pytorch3d/structures/meshes.py | 109 +++++++++++++++++++++++++ tests/test_meshes.py | 140 +++++++++++++++++++++++++++++++++ 2 files changed, 249 insertions(+) diff --git a/pytorch3d/structures/meshes.py b/pytorch3d/structures/meshes.py index 0cc36dcd..7ec5d7e1 100644 --- a/pytorch3d/structures/meshes.py +++ b/pytorch3d/structures/meshes.py @@ -1556,6 +1556,115 @@ class Meshes: else: raise ValueError("Meshes does not have textures") + def submeshes( + self, + face_indices: Union[ + List[List[torch.LongTensor]], List[torch.LongTensor], torch.LongTensor + ], + ) -> "Meshes": + """ + Split a batch of meshes into a batch of submeshes. + + The return value is a Meshes object representing + [mesh restricted to only faces indexed by selected_faces + for mesh, selected_faces_list in zip(self, face_indices) + for faces in selected_faces_list] + + Args: + face_indices: + Let the original mesh have verts_list() of length N. + Can be either + - List of lists of LongTensors. The n-th element is a list of length + num_submeshes_n (empty lists are allowed). The k-th element of the n-th + sublist is a LongTensor of length num_faces_submesh_n_k. + - List of LongTensors. The n-th element is a (possibly empty) LongTensor + of shape (num_submeshes_n, num_faces_n). + - A LongTensor of shape (N, num_submeshes_per_mesh, num_faces_per_submesh) + where all meshes in the batch will have the same number of submeshes. + This will result in an output Meshes object with batch size equal to + N * num_submeshes_per_mesh. + + Returns: + Meshes object of length `sum(len(ids) for ids in face_indices)`. + + Submeshing only works with no textures or with the TexturesVertex texture. + + Example 1: + + If `meshes` has batch size 1, and `face_indices` is a 1D LongTensor, + then `meshes.submeshes([[face_indices]]) and + `meshes.submeshes(face_indices[None, None])` both produce a Meshes of length 1, + containing a single submesh with a subset of `meshes`' faces, whose indices are + specified by `face_indices`. + + Example 2: + + Take a Meshes object `cubes` with 4 meshes, each a translated cube. Then: + * len(cubes) is 4, len(cubes.verts_list()) is 4, len(cubes.faces_list()) 4, + * [cube_verts.size for cube_verts in cubes.verts_list()] is [8, 8, 8, 8], + * [cube_faces.size for cube_faces in cubes.faces_list()] if [6, 6, 6, 6], + + Now let front_facet, top_and_bottom, all_facets be LongTensors of + sizes (2), (4), and (12), each picking up a number of facets of a cube by + specifying the appropriate triangular faces. + + Then let `subcubes = cubes.submeshes([[front_facet, top_and_bottom], [], + [all_facets], []])`. + * len(subcubes) is 3. + * subcubes[0] is the front facet of the cube contained in cubes[0]. + * subcubes[1] is a mesh containing the (disconnected) top and bottom facets + of cubes[0]. + * subcubes[2] is cubes[2]. + * There are no submeshes of cubes[1] and cubes[3] in subcubes. + * subcubes[0] and subcubes[1] are not watertight. subcubes[2] is. + """ + if not ( + self.textures is None or type(self.textures).__name__ == "TexturesVertex" + ): + raise ValueError( + "Submesh extraction only works with no textures or TexturesVertex." + ) + + if len(face_indices) != len(self): + raise ValueError( + "You must specify exactly one set of submeshes" + " for each mesh in this Meshes object." + ) + + sub_verts = [] + sub_faces = [] + + for face_ids_per_mesh, faces, verts in zip( + face_indices, self.faces_list(), self.verts_list() + ): + for submesh_face_ids in face_ids_per_mesh: + faces_to_keep = faces[submesh_face_ids] + + # Say we are keeping two faces from a mesh with six vertices: + # faces_to_keep = [[0, 6, 4], + # [0, 2, 6]] + # Then we want verts_to_keep to contain only vertices [0, 2, 4, 6]: + vertex_ids_to_keep = torch.unique(faces_to_keep, sorted=True) + sub_verts.append(verts[vertex_ids_to_keep]) + + # Now, convert faces_to_keep to use the new vertex ids. + # In our example, instead of + # [[0, 6, 4], + # [0, 2, 6]] + # we want faces_to_keep to be + # [[0, 3, 2], + # [0, 1, 3]], + # as each point id got reduced to its sort rank. + _, ids_of_unique_ids_in_sorted = torch.unique( + faces_to_keep, return_inverse=True + ) + sub_faces.append(ids_of_unique_ids_in_sorted) + + return self.__class__( + verts=sub_verts, + faces=sub_faces, + ) + def join_meshes_as_batch(meshes: List[Meshes], include_textures: bool = True) -> Meshes: """ diff --git a/tests/test_meshes.py b/tests/test_meshes.py index 0b11c054..b8b80ba7 100644 --- a/tests/test_meshes.py +++ b/tests/test_meshes.py @@ -233,6 +233,46 @@ def to_sorted(mesh: Meshes) -> "Meshes": return other +def init_cube_meshes(device: str = "cpu"): + # Make Meshes with four cubes translated from the origin by varying amounts. + verts = torch.FloatTensor( + [ + [0, 0, 0], + [1, 0, 0], # 1->0 + [1, 1, 0], # 2->1 + [0, 1, 0], # 3->2 + [0, 1, 1], # 3 + [1, 1, 1], # 4 + [1, 0, 1], # 5 + [0, 0, 1], + ], + device=device, + ) + + faces = torch.FloatTensor( + [ + [0, 2, 1], + [0, 3, 2], + [2, 3, 4], # 1,2, 3 + [2, 4, 5], # + [1, 2, 5], # + [1, 5, 6], # + [0, 7, 4], + [0, 4, 3], + [5, 4, 7], + [5, 7, 6], + [0, 6, 7], + [0, 1, 6], + ], + device=device, + ) + + return Meshes( + verts=[verts, verts + 1, verts + 2, verts + 3], + faces=[faces, faces, faces, faces], + ) + + class TestMeshes(TestCaseMixin, unittest.TestCase): def setUp(self) -> None: np.random.seed(42) @@ -1257,6 +1297,106 @@ class TestMeshes(TestCaseMixin, unittest.TestCase): yes_normals.offset_verts_(torch.FloatTensor([1, 2, 3]).expand(12, 3)) self.assertFalse(torch.allclose(yes_normals.verts_normals_padded(), verts)) + def test_submeshes(self): + empty_mesh = Meshes([], []) + # Four cubes with offsets [0, 1, 2, 3]. + cubes = init_cube_meshes() + + # Extracting an empty submesh from an empty mesh is allowed, but extracting + # a nonempty submesh from an empty mesh should result in a value error. + self.assertTrue(mesh_structures_equal(empty_mesh.submeshes([]), empty_mesh)) + self.assertTrue( + mesh_structures_equal(cubes.submeshes([[], [], [], []]), empty_mesh) + ) + + with self.assertRaisesRegex( + ValueError, "You must specify exactly one set of submeshes" + ): + empty_mesh.submeshes([torch.LongTensor([0])]) + + # Check that we can chop the cube up into its facets. + subcubes = to_sorted( + cubes.submeshes( + [ # Do not submesh cube#1. + [], + # Submesh the front face and the top-and-bottom of cube#2. + [ + torch.LongTensor([0, 1]), + torch.LongTensor([2, 3, 4, 5]), + ], + # Do not submesh cube#3. + [], + # Submesh the whole cube#4 (clone it). + [torch.LongTensor(list(range(12)))], + ] + ) + ) + + # The cube should've been chopped into three submeshes. + self.assertEquals(len(subcubes), 3) + + # The first submesh should be a single facet of cube#2. + front_facet = to_sorted( + Meshes( + verts=torch.FloatTensor([[[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]]]) + + 1, + faces=torch.LongTensor([[[0, 2, 1], [0, 3, 2]]]), + ) + ) + self.assertTrue(mesh_structures_equal(front_facet, subcubes[0])) + + # The second submesh should be the top and bottom facets of cube#2. + top_and_bottom = Meshes( + verts=torch.FloatTensor( + [[[1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 1, 1], [1, 1, 1], [1, 0, 1]]] + ) + + 1, + faces=torch.LongTensor([[[1, 2, 3], [1, 3, 4], [0, 1, 4], [0, 4, 5]]]), + ) + self.assertTrue(mesh_structures_equal(to_sorted(top_and_bottom), subcubes[1])) + + # The last submesh should be all of cube#3. + self.assertTrue(mesh_structures_equal(to_sorted(cubes[3]), subcubes[2])) + + # Test alternative input parameterization: list of LongTensors. + two_facets = torch.LongTensor([[0, 1], [4, 5]]) + subcubes = to_sorted(cubes.submeshes([two_facets, [], two_facets, []])) + expected_verts = torch.FloatTensor( + [ + [[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 1, 0]], + [[1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]], + [[2, 2, 2], [2, 3, 2], [3, 2, 2], [3, 3, 2]], + [[3, 2, 2], [3, 2, 3], [3, 3, 2], [3, 3, 3]], + ] + ) + expected_faces = torch.LongTensor( + [ + [[0, 3, 2], [0, 1, 3]], + [[0, 2, 3], [0, 3, 1]], + [[0, 3, 2], [0, 1, 3]], + [[0, 2, 3], [0, 3, 1]], + ] + ) + expected_meshes = Meshes(verts=expected_verts, faces=expected_faces) + self.assertTrue(mesh_structures_equal(subcubes, expected_meshes)) + + # Test alternative input parameterization: a single LongTensor. + triangle_per_mesh = torch.LongTensor([[[0]], [[1]], [[4]], [[5]]]) + subcubes = to_sorted(cubes.submeshes(triangle_per_mesh)) + expected_verts = torch.FloatTensor( + [ + [[0, 0, 0], [1, 0, 0], [1, 1, 0]], + [[1, 1, 1], [1, 2, 1], [2, 2, 1]], + [[3, 2, 2], [3, 3, 2], [3, 3, 3]], + [[4, 3, 3], [4, 3, 4], [4, 4, 4]], + ] + ) + expected_faces = torch.LongTensor( + [[[0, 2, 1]], [[0, 1, 2]], [[0, 1, 2]], [[0, 2, 1]]] + ) + expected_meshes = Meshes(verts=expected_verts, faces=expected_faces) + self.assertTrue(mesh_structures_equal(subcubes, expected_meshes)) + def test_compute_faces_areas_cpu_cuda(self): num_meshes = 10 max_v = 100