Submesh 4/n: TexturesVertex submeshing

Krzysztof Chalupka · facebook-github-bot · commit 22f86072ca1c · 2022-04-11T16:27:53.000-07:00
Summary: Add submeshing capability for meshes with TexturesVertex.

Reviewed By: bottler

Differential Revision: D35448534

fbshipit-source-id: 6d16a31a5bfb24ce122cf3c300a7616bc58353d1
diff --git a/pytorch3d/renderer/mesh/textures.py b/pytorch3d/renderer/mesh/textures.py
@@ -242,6 +242,16 @@ def sample_textures(self) -> torch.Tensor:
         """
         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 @@ def sample_textures(self, fragments, faces_packed=None) -> torch.Tensor:
         )
         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
@@ -1618,34 +1618,33 @@ def submeshes(
             * 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"
                 " for each mesh in this Meshes object."
             )
 
         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],
                 #                  [0, 2, 6]]
                 # 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 @@ def submeshes(
         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
@@ -145,6 +145,53 @@ def test_faces_verts_textures(self):
 
         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()
