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