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

tests/test_texturing.py

@@ -576,6 +576,39 @@ class TestTexturesAtlas(TestCaseMixin, unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
             meshes.sample_textures(None)
 
+    def test_submeshes(self):
+        N = 2
+        V = 5
+        F = 5
+        tex = TexturesAtlas(
+            atlas=torch.arange(N * F * 4 * 4 * 3, dtype=torch.float32).reshape(
+                N, F, 4, 4, 3
+            )
+        )
+
+        verts = torch.rand(size=(N, V, 3))
+        faces = torch.randint(size=(N, F, 3), high=V)
+        mesh = Meshes(verts=verts, faces=faces, textures=tex)
+
+        sub_faces = [
+            [torch.tensor([0, 2]), torch.tensor([1, 2])],
+            [],
+        ]
+        subtex = mesh.submeshes(sub_faces).textures
+        subtex_faces = subtex.atlas_list()
+
+        self.assertEqual(len(subtex_faces), 2)
+        self.assertClose(
+            subtex_faces[0].flatten().msort(),
+            torch.cat(
+                (
+                    torch.arange(4 * 4 * 3, dtype=torch.float32),
+                    torch.arange(96, 96 + 4 * 4 * 3, dtype=torch.float32),
+                ),
+                0,
+            ),
+        )
+
 
 class TestTexturesUV(TestCaseMixin, unittest.TestCase):
     def setUp(self) -> None: