From 22f86072ca1cda37a679fc7937311af47bf8fa3b Mon Sep 17 00:00:00 2001 From: Krzysztof Chalupka Date: Mon, 11 Apr 2022 16:27:53 -0700 Subject: [PATCH] Submesh 4/n: TexturesVertex submeshing Summary: Add submeshing capability for meshes with TexturesVertex. Reviewed By: bottler Differential Revision: D35448534 fbshipit-source-id: 6d16a31a5bfb24ce122cf3c300a7616bc58353d1 --- pytorch3d/renderer/mesh/textures.py | 51 +++++++++++++++++++++++++++++ pytorch3d/structures/meshes.py | 18 ++++++---- tests/test_texturing.py | 47 ++++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 7 deletions(-) diff --git a/pytorch3d/renderer/mesh/textures.py b/pytorch3d/renderer/mesh/textures.py index d0bbe82f..3a83dd7f 100644 --- a/pytorch3d/renderer/mesh/textures.py +++ b/pytorch3d/renderer/mesh/textures.py @@ -242,6 +242,16 @@ class TexturesBase: """ raise NotImplementedError() + def submeshes( + self, + vertex_ids_list: Optional[List[List[torch.LongTensor]]], + faces_ids_list: Optional[List[List[torch.LongTensor]]], + ) -> "TexturesBase": + """ + Extract sub-textures used for submeshing. + """ + raise NotImplementedError(f"{self.__class__} does not support submeshes") + def faces_verts_textures_packed(self) -> torch.Tensor: """ Returns the texture for each vertex for each face in the mesh. @@ -1465,6 +1475,47 @@ class TexturesVertex(TexturesBase): ) return texels + def submeshes( + self, + vertex_ids_list: Optional[List[List[torch.LongTensor]]], + faces_ids_list: Optional[List[List[torch.LongTensor]]], + ) -> "TexturesVertex": + """ + Extract a sub-texture for use in a submesh. + + If the meshes batch corresponding to this TexturesVertex contains + `n = len(vertex_ids_list)` meshes, then self.verts_features_list() + will be of length n. After submeshing, we obtain a batch of + `k = sum(len(v) for v in vertex_ids_list` submeshes (see Meshes.submeshes). This + function creates a corresponding TexturesVertex object with `verts_features_list` + of length `k`. + + Args: + vertex_ids_list: A list of length equal to self.verts_features_list. Each + element is a LongTensor listing the vertices that the submesh keeps in + each respective mesh. + + face_ids_list: Not used when submeshing TexturesVertex. + + Returns: + A TexturesVertex in which verts_features_list has length + sum(len(vertices) for vertices in vertex_ids_list). Each element contains + vertex features corresponding to the subset of vertices in that submesh. + """ + if vertex_ids_list is None or len(vertex_ids_list) != len( + self.verts_features_list() + ): + raise IndexError( + "verts_features_list must be of " "the same length as vertex_ids_list." + ) + + sub_features = [] + for vertex_ids, features in zip(vertex_ids_list, self.verts_features_list()): + for vertex_ids_submesh in vertex_ids: + sub_features.append(features[vertex_ids_submesh]) + + return self.__class__(sub_features) + def faces_verts_textures_packed(self, faces_packed=None) -> torch.Tensor: """ Samples texture from each vertex and for each face in the mesh. diff --git a/pytorch3d/structures/meshes.py b/pytorch3d/structures/meshes.py index 7ec5d7e1..1761e7bf 100644 --- a/pytorch3d/structures/meshes.py +++ b/pytorch3d/structures/meshes.py @@ -1618,13 +1618,6 @@ class Meshes: * 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" @@ -1632,13 +1625,18 @@ class Meshes: ) sub_verts = [] + sub_verts_ids = [] sub_faces = [] + sub_face_ids = [] for face_ids_per_mesh, faces, verts in zip( face_indices, self.faces_list(), self.verts_list() ): + sub_verts_ids.append([]) + sub_face_ids.append([]) for submesh_face_ids in face_ids_per_mesh: faces_to_keep = faces[submesh_face_ids] + sub_face_ids[-1].append(faces_to_keep) # Say we are keeping two faces from a mesh with six vertices: # faces_to_keep = [[0, 6, 4], @@ -1646,6 +1644,7 @@ class Meshes: # 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]) + sub_verts_ids[-1].append(vertex_ids_to_keep) # Now, convert faces_to_keep to use the new vertex ids. # In our example, instead of @@ -1663,6 +1662,11 @@ class Meshes: return self.__class__( verts=sub_verts, faces=sub_faces, + textures=( + self.textures.submeshes(sub_verts_ids, sub_face_ids) + if self.textures + else None + ), ) diff --git a/tests/test_texturing.py b/tests/test_texturing.py index e4c5286f..a4e8cf15 100644 --- a/tests/test_texturing.py +++ b/tests/test_texturing.py @@ -145,6 +145,53 @@ class TestTexturesVertex(TestCaseMixin, unittest.TestCase): self.assertClose(faces_verts_texts_packed, faces_verts_texts) + def test_submeshes(self): + # define TexturesVertex + verts_features = torch.tensor( + [ + [1, 0, 0], + [1, 0, 0], + [1, 0, 0], + [1, 0, 0], + [0, 1, 0], + [0, 1, 0], + [0, 1, 0], + [0, 1, 0], + ], + dtype=torch.float32, + ) + + textures = TexturesVertex( + verts_features=[verts_features, verts_features, verts_features] + ) + subtextures = textures.submeshes( + [ + [ + torch.LongTensor([0, 2, 3]), + torch.LongTensor(list(range(8))), + ], + [], + [ + torch.LongTensor([4]), + ], + ], + None, + ) + + subtextures_features = subtextures.verts_features_list() + + self.assertEqual(len(subtextures_features), 3) + self.assertTrue( + torch.equal( + subtextures_features[0], + torch.FloatTensor([[1, 0, 0], [1, 0, 0], [1, 0, 0]]), + ) + ) + self.assertTrue(torch.equal(subtextures_features[1], verts_features)) + self.assertTrue( + torch.equal(subtextures_features[2], torch.FloatTensor([[0, 1, 0]])) + ) + def test_clone(self): tex = TexturesVertex(verts_features=torch.rand(size=(10, 100, 128))) tex.verts_features_list()