Submesh 3/n: Add submeshing functionality

Summary:
Copypasting the docstring:
```
        Split a mesh into submeshes, defined by face indices of the original Meshes object.

        Args:
          face_indices:
            Let the original mesh have verts_list() of length N.
            Can be either
              - List of length N. The n-th element is a list of length num_submeshes_n
                (empty lists are allowed). Each element of the n-th sublist is a LongTensor
                of length num_faces.
              - List of length N. The n-th element is a possibly empty padded LongTensor of
                shape (num_submeshes_n, max_num_faces).

        Returns:
          Meshes object with selected submeshes. The submesh tensors are cloned.

        Currently submeshing only works with no textures or with the TexturesVertex texture.

        Example:

        Take a Meshes object `cubes` with 4 meshes, each a translated cube. Then:
            * len(cubes) is 4, len(cubes.verts_list()) is 4, len(cubes.faces_list()) is 4,
            * [cube_verts.size for cube_verts in cubes.verts_list()] is [8, 8, 8, 8],
            * [cube_faces.size for cube_faces in cubes.faces_list()] if [6, 6, 6, 6],

        Now let front_facet, top_and_bottom, all_facets be LongTensors of
        sizes (2), (4), and (12), each picking up a number of facets of a cube by specifying
        the appropriate triangular faces.

        Then let `subcubes = cubes.submeshes([[front_facet, top_and_bottom], [], [all_facets], []])`.
            * len(subcubes) is 3.
            * subcubes[0] is the front facet of the cube contained in cubes[0].
            * subcubes[1] is a mesh containing the (disconnected) top and bottom facets of cubes[0].
            * subcubes[2] is a clone of cubes[2].
            * There are no submeshes of cubes[1] and cubes[3] in subcubes.
            * subcubes[0] and subcubes[1] are not watertight. subcubes[2] is.
```

Reviewed By: bottler

Differential Revision: D35440657

fbshipit-source-id: 8a6d2d300ce226b5b9eb440688528b5e795195a1

pytorch3d/structures/meshes.py

@@ -1556,6 +1556,115 @@ def sample_textures(self, fragments):
         else:
             raise ValueError("Meshes does not have textures")
 
+    def submeshes(
+        self,
+        face_indices: Union[
+            List[List[torch.LongTensor]], List[torch.LongTensor], torch.LongTensor
+        ],
+    ) -> "Meshes":
+        """
+        Split a batch of meshes into a batch of submeshes.
+
+        The return value is a Meshes object representing
+            [mesh restricted to only faces indexed by selected_faces
+            for mesh, selected_faces_list in zip(self, face_indices)
+            for faces in selected_faces_list]
+
+        Args:
+          face_indices:
+            Let the original mesh have verts_list() of length N.
+            Can be either
+              - List of lists of LongTensors. The n-th element is a list of length
+              num_submeshes_n (empty lists are allowed). The k-th element of the n-th
+              sublist is a LongTensor of length num_faces_submesh_n_k.
+              - List of LongTensors. The n-th element is a (possibly empty) LongTensor
+                of shape (num_submeshes_n, num_faces_n).
+              - A LongTensor of shape (N, num_submeshes_per_mesh, num_faces_per_submesh)
+                where all meshes in the batch will have the same number of submeshes.
+                This will result in an output Meshes object with batch size equal to
+                N * num_submeshes_per_mesh.
+
+        Returns:
+          Meshes object of length `sum(len(ids) for ids in face_indices)`.
+
+        Submeshing only works with no textures or with the TexturesVertex texture.
+
+        Example 1:
+
+        If `meshes` has batch size 1, and `face_indices` is a 1D LongTensor,
+        then `meshes.submeshes([[face_indices]]) and
+        `meshes.submeshes(face_indices[None, None])` both produce a Meshes of length 1,
+        containing a single submesh with a subset of `meshes`' faces, whose indices are
+        specified by `face_indices`.
+
+        Example 2:
+
+        Take a Meshes object `cubes` with 4 meshes, each a translated cube. Then:
+            * len(cubes) is 4, len(cubes.verts_list()) is 4, len(cubes.faces_list()) 4,
+            * [cube_verts.size for cube_verts in cubes.verts_list()] is [8, 8, 8, 8],
+            * [cube_faces.size for cube_faces in cubes.faces_list()] if [6, 6, 6, 6],
+
+        Now let front_facet, top_and_bottom, all_facets be LongTensors of
+        sizes (2), (4), and (12), each picking up a number of facets of a cube by
+        specifying the appropriate triangular faces.
+
+        Then let `subcubes = cubes.submeshes([[front_facet, top_and_bottom], [],
+                                              [all_facets], []])`.
+            * len(subcubes) is 3.
+            * subcubes[0] is the front facet of the cube contained in cubes[0].
+            * subcubes[1] is a mesh containing the (disconnected) top and bottom facets
+              of cubes[0].
+            * subcubes[2] is cubes[2].
+            * 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_faces = []
+
+        for face_ids_per_mesh, faces, verts in zip(
+            face_indices, self.faces_list(), self.verts_list()
+        ):
+            for submesh_face_ids in face_ids_per_mesh:
+                faces_to_keep = faces[submesh_face_ids]
+
+                # 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])
+
+                # Now, convert faces_to_keep to use the new vertex ids.
+                # In our example, instead of
+                # [[0, 6, 4],
+                #  [0, 2, 6]]
+                # we want faces_to_keep to be
+                # [[0, 3, 2],
+                #  [0, 1, 3]],
+                # as each point id got reduced to its sort rank.
+                _, ids_of_unique_ids_in_sorted = torch.unique(
+                    faces_to_keep, return_inverse=True
+                )
+                sub_faces.append(ids_of_unique_ids_in_sorted)
+
+        return self.__class__(
+            verts=sub_verts,
+            faces=sub_faces,
+        )
+
 
 def join_meshes_as_batch(meshes: List[Meshes], include_textures: bool = True) -> Meshes:
     """

tests/test_meshes.py

@@ -233,6 +233,46 @@ def to_sorted(mesh: Meshes) -> "Meshes":
     return other
 
 
+def init_cube_meshes(device: str = "cpu"):
+    # Make Meshes with four cubes translated from the origin by varying amounts.
+    verts = torch.FloatTensor(
+        [
+            [0, 0, 0],
+            [1, 0, 0],  # 1->0
+            [1, 1, 0],  # 2->1
+            [0, 1, 0],  # 3->2
+            [0, 1, 1],  # 3
+            [1, 1, 1],  # 4
+            [1, 0, 1],  # 5
+            [0, 0, 1],
+        ],
+        device=device,
+    )
+
+    faces = torch.FloatTensor(
+        [
+            [0, 2, 1],
+            [0, 3, 2],
+            [2, 3, 4],  # 1,2, 3
+            [2, 4, 5],  #
+            [1, 2, 5],  #
+            [1, 5, 6],  #
+            [0, 7, 4],
+            [0, 4, 3],
+            [5, 4, 7],
+            [5, 7, 6],
+            [0, 6, 7],
+            [0, 1, 6],
+        ],
+        device=device,
+    )
+
+    return Meshes(
+        verts=[verts, verts + 1, verts + 2, verts + 3],
+        faces=[faces, faces, faces, faces],
+    )
+
+
 class TestMeshes(TestCaseMixin, unittest.TestCase):
     def setUp(self) -> None:
         np.random.seed(42)
@@ -1257,6 +1297,106 @@ def test_assigned_normals(self):
             yes_normals.offset_verts_(torch.FloatTensor([1, 2, 3]).expand(12, 3))
             self.assertFalse(torch.allclose(yes_normals.verts_normals_padded(), verts))
 
+    def test_submeshes(self):
+        empty_mesh = Meshes([], [])
+        # Four cubes with offsets [0, 1, 2, 3].
+        cubes = init_cube_meshes()
+
+        # Extracting an empty submesh from an empty mesh is allowed, but extracting
+        # a nonempty submesh from an empty mesh should result in a value error.
+        self.assertTrue(mesh_structures_equal(empty_mesh.submeshes([]), empty_mesh))
+        self.assertTrue(
+            mesh_structures_equal(cubes.submeshes([[], [], [], []]), empty_mesh)
+        )
+
+        with self.assertRaisesRegex(
+            ValueError, "You must specify exactly one set of submeshes"
+        ):
+            empty_mesh.submeshes([torch.LongTensor([0])])
+
+        # Check that we can chop the cube up into its facets.
+        subcubes = to_sorted(
+            cubes.submeshes(
+                [  # Do not submesh cube#1.
+                    [],
+                    # Submesh the front face and the top-and-bottom of cube#2.
+                    [
+                        torch.LongTensor([0, 1]),
+                        torch.LongTensor([2, 3, 4, 5]),
+                    ],
+                    # Do not submesh cube#3.
+                    [],
+                    # Submesh the whole cube#4 (clone it).
+                    [torch.LongTensor(list(range(12)))],
+                ]
+            )
+        )
+
+        # The cube should've been chopped into three submeshes.
+        self.assertEquals(len(subcubes), 3)
+
+        # The first submesh should be a single facet of cube#2.
+        front_facet = to_sorted(
+            Meshes(
+                verts=torch.FloatTensor([[[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]]])
+                + 1,
+                faces=torch.LongTensor([[[0, 2, 1], [0, 3, 2]]]),
+            )
+        )
+        self.assertTrue(mesh_structures_equal(front_facet, subcubes[0]))
+
+        # The second submesh should be the top and bottom facets of cube#2.
+        top_and_bottom = Meshes(
+            verts=torch.FloatTensor(
+                [[[1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 1, 1], [1, 1, 1], [1, 0, 1]]]
+            )
+            + 1,
+            faces=torch.LongTensor([[[1, 2, 3], [1, 3, 4], [0, 1, 4], [0, 4, 5]]]),
+        )
+        self.assertTrue(mesh_structures_equal(to_sorted(top_and_bottom), subcubes[1]))
+
+        # The last submesh should be all of cube#3.
+        self.assertTrue(mesh_structures_equal(to_sorted(cubes[3]), subcubes[2]))
+
+        # Test alternative input parameterization: list of LongTensors.
+        two_facets = torch.LongTensor([[0, 1], [4, 5]])
+        subcubes = to_sorted(cubes.submeshes([two_facets, [], two_facets, []]))
+        expected_verts = torch.FloatTensor(
+            [
+                [[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 1, 0]],
+                [[1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]],
+                [[2, 2, 2], [2, 3, 2], [3, 2, 2], [3, 3, 2]],
+                [[3, 2, 2], [3, 2, 3], [3, 3, 2], [3, 3, 3]],
+            ]
+        )
+        expected_faces = torch.LongTensor(
+            [
+                [[0, 3, 2], [0, 1, 3]],
+                [[0, 2, 3], [0, 3, 1]],
+                [[0, 3, 2], [0, 1, 3]],
+                [[0, 2, 3], [0, 3, 1]],
+            ]
+        )
+        expected_meshes = Meshes(verts=expected_verts, faces=expected_faces)
+        self.assertTrue(mesh_structures_equal(subcubes, expected_meshes))
+
+        # Test alternative input parameterization: a single LongTensor.
+        triangle_per_mesh = torch.LongTensor([[[0]], [[1]], [[4]], [[5]]])
+        subcubes = to_sorted(cubes.submeshes(triangle_per_mesh))
+        expected_verts = torch.FloatTensor(
+            [
+                [[0, 0, 0], [1, 0, 0], [1, 1, 0]],
+                [[1, 1, 1], [1, 2, 1], [2, 2, 1]],
+                [[3, 2, 2], [3, 3, 2], [3, 3, 3]],
+                [[4, 3, 3], [4, 3, 4], [4, 4, 4]],
+            ]
+        )
+        expected_faces = torch.LongTensor(
+            [[[0, 2, 1]], [[0, 1, 2]], [[0, 1, 2]], [[0, 2, 1]]]
+        )
+        expected_meshes = Meshes(verts=expected_verts, faces=expected_faces)
+        self.assertTrue(mesh_structures_equal(subcubes, expected_meshes))
+
     def test_compute_faces_areas_cpu_cuda(self):
         num_meshes = 10
         max_v = 100