diff --git a/pytorch3d/renderer/mesh/textures.py b/pytorch3d/renderer/mesh/textures.py index 5f134635..e38cf7c0 100644 --- a/pytorch3d/renderer/mesh/textures.py +++ b/pytorch3d/renderer/mesh/textures.py @@ -1332,6 +1332,60 @@ class TexturesUV(TexturesBase): self.verts_uvs_padded().shape[0] == batch_size ) + def submeshes( + self, + vertex_ids_list: List[List[torch.LongTensor]], + faces_ids_list: List[List[torch.LongTensor]], + ) -> "TexturesUV": + """ + Extract a sub-texture for use in a submesh. + + If the meshes batch corresponding to this TexturesUV contains + `n = len(faces_ids_list)` meshes, then self.faces_uvs_padded() + will be of length n. After submeshing, we obtain a batch of + `k = sum(len(f) for f in faces_ids_list` submeshes (see Meshes.submeshes). This + function creates a corresponding TexturesUV object with `faces_uvs_padded` + of length `k`. + + Args: + vertex_ids_list: Not used when submeshing TexturesUV. + + face_ids_list: A list of length equal to self.faces_uvs_padded. Each + element is a LongTensor listing the face ids that the submesh keeps in + each respective mesh. + + + Returns: + A "TexturesUV in which faces_uvs_padded, verts_uvs_padded, and maps_padded + have length sum(len(faces) for faces in faces_ids_list) + """ + + if len(faces_ids_list) != len(self.faces_uvs_padded()): + raise IndexError( + "faces_uvs_padded must be of " "the same length as face_ids_list." + ) + + sub_faces_uvs, sub_verts_uvs, sub_maps = [], [], [] + for faces_ids, faces_uvs, verts_uvs, map_ in zip( + faces_ids_list, + self.faces_uvs_padded(), + self.verts_uvs_padded(), + self.maps_padded(), + ): + for faces_ids_submesh in faces_ids: + sub_faces_uvs.append(faces_uvs[faces_ids_submesh]) + sub_verts_uvs.append(verts_uvs) + sub_maps.append(map_) + + return self.__class__( + sub_maps, + sub_faces_uvs, + sub_verts_uvs, + self.padding_mode, + self.align_corners, + self.sampling_mode, + ) + class TexturesVertex(TexturesBase): def __init__( diff --git a/pytorch3d/structures/meshes.py b/pytorch3d/structures/meshes.py index e1ad35d2..19b8fa5d 100644 --- a/pytorch3d/structures/meshes.py +++ b/pytorch3d/structures/meshes.py @@ -1576,7 +1576,7 @@ class Meshes: Returns: Meshes object of length `sum(len(ids) for ids in face_indices)`. - Submeshing only works with no textures or with the TexturesVertex texture. + Submeshing only works with no textures, TexturesVertex, or TexturesUV. Example 1: @@ -1616,16 +1616,13 @@ 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], @@ -1652,7 +1649,7 @@ class Meshes: verts=sub_verts, faces=sub_faces, textures=( - self.textures.submeshes(sub_verts_ids, sub_face_ids) + self.textures.submeshes(sub_verts_ids, face_indices) if self.textures else None ), diff --git a/tests/test_texturing.py b/tests/test_texturing.py index 35269586..5234f2ae 100644 --- a/tests/test_texturing.py +++ b/tests/test_texturing.py @@ -1002,6 +1002,49 @@ class TestTexturesUV(TestCaseMixin, unittest.TestCase): with self.assertRaisesRegex(ValueError, "do not match the dimensions"): meshes.sample_textures(None) + def test_submeshes(self): + N = 2 + faces_uvs_list = [ + torch.LongTensor([[0, 1, 2], [3, 5, 4], [7, 6, 8]]), + torch.LongTensor([[0, 1, 2], [3, 4, 5]]), + ] + verts_uvs_list = [ + torch.arange(18, dtype=torch.float32).reshape(9, 2), + torch.ones(6, 2), + ] + tex = TexturesUV( + maps=torch.rand((N, 16, 16, 3)), + faces_uvs=faces_uvs_list, + verts_uvs=verts_uvs_list, + ) + + sub_faces = [ + [torch.tensor([0, 1]), torch.tensor([1, 2])], + [], + ] + + mesh = Meshes( + verts=[torch.rand(9, 3), torch.rand(6, 3)], + faces=faces_uvs_list, + textures=tex, + ) + subtex = mesh.submeshes(sub_faces).textures + subtex_faces = subtex.faces_uvs_padded() + self.assertEqual(len(subtex_faces), 2) + self.assertClose( + subtex_faces[0], + torch.tensor([[0, 1, 2], [3, 5, 4]]), + ) + self.assertClose( + subtex.verts_uvs_list()[0][subtex.faces_uvs_list()[0].flatten()] + .flatten() + .msort(), + torch.arange(12, dtype=torch.float32), + ) + self.assertClose( + subtex.maps_padded(), tex.maps_padded()[:1].expand(2, -1, -1, -1) + ) + class TestRectanglePacking(TestCaseMixin, unittest.TestCase): def setUp(self) -> None: