Submeshing TexturesUV

Summary: Implement `submeshes` for TexturesUV. Fix what Meshes.submeshes passes to the texture's submeshes function to make this possible.

Reviewed By: bottler

Differential Revision: D52192060

fbshipit-source-id: 526734962e3376aaf75654200164cdcebfff6997

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__(

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

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: