Submeshing TexturesAtlas for PyTorch3D 3D Rendering

Summary: Implement submeshing for TexturesAtlas and add associated test

Reviewed By: bottler

Differential Revision: D52334053

fbshipit-source-id: d54080e9af1f0c01551702736e858e3bd439ac58

pytorch3d/renderer/mesh/textures.py

@@ -549,6 +549,33 @@ class TexturesAtlas(TexturesBase):
 
         return texels
 
+    def submeshes(
+        self,
+        vertex_ids_list: List[List[torch.LongTensor]],
+        faces_ids_list: List[List[torch.LongTensor]],
+    ) -> "TexturesAtlas":
+        """
+        Extract a sub-texture for use in a submesh.
+
+        If the meshes batch corresponding to this TextureAtlas contains
+        `n = len(faces_ids_list)` meshes, then self.atlas_list()
+        will be of length n. After submeshing, we obtain a batch of
+        `k = sum(len(v) for v in atlas_list` submeshes (see Meshes.submeshes). This
+        function creates a corresponding TexturesAtlas object with `atlas_list`
+        of length `k`.
+        """
+        if len(faces_ids_list) != len(self.atlas_list()):
+            raise IndexError(
+                "faces_ids_list must be of " "the same length as atlas_list."
+            )
+
+        sub_features = []
+        for atlas, faces_ids in zip(self.atlas_list(), faces_ids_list):
+            for faces_ids_submesh in faces_ids:
+                sub_features.append(atlas[faces_ids_submesh])
+
+        return self.__class__(sub_features)
+
     def faces_verts_textures_packed(self) -> torch.Tensor:
         """
         Samples texture from each vertex for each face in the mesh.

pytorch3d/structures/meshes.py

@@ -1576,8 +1576,6 @@ class Meshes:
         Returns:
           Meshes object of length `sum(len(ids) for ids in face_indices)`.
 
-        Submeshing only works with no textures, TexturesVertex, or TexturesUV.
-
         Example 1:
 
         If `meshes` has batch size 1, and `face_indices` is a 1D LongTensor,