Write meshes to GLB

Summary: Write the amalgamated mesh from the Mesh module to glb. In this version, the json header and the binary data specified by the buffer are merged into glb. The image texture attributes are added.

Reviewed By: bottler

Differential Revision: D41489778

fbshipit-source-id: 3af0e9a8f9e9098e73737a254177802e0fb6bd3c

tests/test_io_gltf.py

@@ -29,6 +29,7 @@ from pytorch3d.renderer.mesh import (
 )
 from pytorch3d.structures import Meshes
 from pytorch3d.transforms import axis_angle_to_matrix
+from pytorch3d.utils import ico_sphere
 from pytorch3d.vis.texture_vis import texturesuv_image_PIL
 
 from .common_testing import get_pytorch3d_dir, get_tests_dir, TestCaseMixin
@@ -45,6 +46,12 @@ def _load(path, **kwargs) -> Meshes:
     return io.load_mesh(path, **kwargs)
 
 
+def _write(mesh, path, **kwargs) -> bool:
+    io = IO()
+    io.register_meshes_format(MeshGlbFormat())
+    return io.save_mesh(mesh, path, **kwargs)
+
+
 def _render(
     mesh: Meshes,
     name: str,
@@ -144,9 +151,7 @@ class TestMeshGltfIO(TestCaseMixin, unittest.TestCase):
         self.assertEqual(mesh.faces_packed().shape, (5856, 3))
         self.assertEqual(mesh.verts_packed().shape, (3225, 3))
         mesh_obj = _load(TUTORIAL_DATA_DIR / "cow_mesh/cow.obj")
-        self.assertClose(
-            mesh_obj.get_bounding_boxes().cpu(), mesh_obj.get_bounding_boxes()
-        )
+        self.assertClose(mesh.get_bounding_boxes().cpu(), mesh_obj.get_bounding_boxes())
 
         self.assertClose(
             mesh.textures.verts_uvs_padded().cpu(), mesh_obj.textures.verts_uvs_padded()
@@ -169,6 +174,69 @@ class TestMeshGltfIO(TestCaseMixin, unittest.TestCase):
 
             self.assertClose(image, expected)
 
+    def test_save_cow(self):
+        """
+        Save the cow mesh to a glb file
+        """
+        # load cow mesh from a glb file
+        glb = DATA_DIR / "cow.glb"
+        self.assertTrue(glb.is_file())
+        device = torch.device("cuda:0")
+        mesh = _load(glb, device=device)
+
+        # save the mesh to a glb file
+        glb = DATA_DIR / "cow_write.glb"
+        _write(mesh, glb)
+
+        # load again
+        glb_reload = DATA_DIR / "cow_write.glb"
+        self.assertTrue(glb_reload.is_file())
+        device = torch.device("cuda:0")
+        mesh_reload = _load(glb_reload, device=device)
+
+        # assertions
+        self.assertEqual(mesh_reload.faces_packed().shape, (5856, 3))
+        self.assertEqual(mesh_reload.verts_packed().shape, (3225, 3))
+        self.assertClose(
+            mesh_reload.get_bounding_boxes().cpu(), mesh.get_bounding_boxes().cpu()
+        )
+
+        self.assertClose(
+            mesh_reload.textures.verts_uvs_padded().cpu(),
+            mesh.textures.verts_uvs_padded().cpu(),
+        )
+
+        self.assertClose(
+            mesh_reload.textures.faces_uvs_padded().cpu(),
+            mesh.textures.faces_uvs_padded().cpu(),
+        )
+
+        self.assertClose(
+            mesh_reload.textures.maps_padded().cpu(), mesh.textures.maps_padded().cpu()
+        )
+
+    def test_save_ico_sphere(self):
+        """
+        save the ico_sphere mesh in a glb file
+        """
+        ico_sphere_mesh = ico_sphere(level=3)
+        glb = DATA_DIR / "ico_sphere.glb"
+        _write(ico_sphere_mesh, glb)
+
+        # reload the ico_sphere
+        device = torch.device("cuda:0")
+        mesh_reload = _load(glb, device=device, include_textures=False)
+
+        self.assertClose(
+            ico_sphere_mesh.verts_padded().cpu(),
+            mesh_reload.verts_padded().cpu(),
+        )
+
+        self.assertClose(
+            ico_sphere_mesh.faces_padded().cpu(),
+            mesh_reload.faces_padded().cpu(),
+        )
+
     def test_load_cow_no_texture(self):
         """
         Load the cow as converted to a single mesh in a glb file.
@@ -183,9 +251,7 @@ class TestMeshGltfIO(TestCaseMixin, unittest.TestCase):
         self.assertEqual(mesh.faces_packed().shape, (5856, 3))
         self.assertEqual(mesh.verts_packed().shape, (3225, 3))
         mesh_obj = _load(TUTORIAL_DATA_DIR / "cow_mesh/cow.obj")
-        self.assertClose(
-            mesh_obj.get_bounding_boxes().cpu(), mesh_obj.get_bounding_boxes()
-        )
+        self.assertClose(mesh.get_bounding_boxes().cpu(), mesh_obj.get_bounding_boxes())
 
         mesh.textures = TexturesVertex(0.5 * torch.ones_like(mesh.verts_padded()))