test_io_gltf.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import os.path
import unittest
from math import radians

import numpy as np
import torch
from PIL import Image
from pytorch3d.io import IO
from pytorch3d.io.experimental_gltf_io import _read_header, MeshGlbFormat
from pytorch3d.renderer import (
    AmbientLights,
    BlendParams,
    FoVPerspectiveCameras,
    look_at_view_transform,
    PointLights,
    RasterizationSettings,
    rotate_on_spot,
)
from pytorch3d.renderer.mesh import (
    HardPhongShader,
    MeshRasterizer,
    MeshRenderer,
    TexturesVertex,
)
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


DATA_DIR = get_tests_dir() / "data"
TUTORIAL_DATA_DIR = get_pytorch3d_dir() / "docs/tutorials/data"
DEBUG = False


def _load(path, **kwargs) -> Meshes:
    io = IO()
    io.register_meshes_format(MeshGlbFormat())
    return io.load_mesh(path, **kwargs)


def _write(mesh, path, **kwargs) -> None:
    io = IO()
    io.register_meshes_format(MeshGlbFormat())
    io.save_mesh(mesh, path, **kwargs)

    with open(path, "rb") as f:
        _, stored_length = _read_header(f)
    assert stored_length == os.path.getsize(path)


def _render(
    mesh: Meshes,
    name: str,
    dist: float = 3.0,
    elev: float = 10.0,
    azim: float = 0,
    image_size: int = 256,
    pan=None,
    RT=None,
    use_ambient=False,
):
    device = mesh.device
    if RT is not None:
        R, T = RT
    else:
        R, T = look_at_view_transform(dist, elev, azim)
        if pan is not None:
            R, T = rotate_on_spot(R, T, pan)
    cameras = FoVPerspectiveCameras(device=device, R=R, T=T)

    raster_settings = RasterizationSettings(
        image_size=image_size, blur_radius=0.0, faces_per_pixel=1
    )

    # Init shader settings
    if use_ambient:
        lights = AmbientLights(device=device)
    else:
        lights = PointLights(device=device)
        lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None]

    blend_params = BlendParams(
        sigma=1e-1,
        gamma=1e-4,
        background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
    )
    # Init renderer
    renderer = MeshRenderer(
        rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
        shader=HardPhongShader(
            device=device, lights=lights, cameras=cameras, blend_params=blend_params
        ),
    )

    output = renderer(mesh)

    image = (output[0, ..., :3].cpu().numpy() * 255).astype(np.uint8)

    if DEBUG:
        Image.fromarray(image).save(DATA_DIR / f"glb_{name}_.png")

    return image


class TestMeshGltfIO(TestCaseMixin, unittest.TestCase):
    def test_load_apartment(self):
        """
        This is the example habitat example scene from inside
        http://dl.fbaipublicfiles.com/habitat/habitat-test-scenes.zip

        The scene is "already lit", i.e. the textures reflect the lighting
        already, so we want to render them with full ambient light.
        """

        self.skipTest("Data not available")

        glb = DATA_DIR / "apartment_1.glb"
        self.assertTrue(glb.is_file())
        device = torch.device("cuda:0")
        mesh = _load(glb, device=device)

        if DEBUG:
            texturesuv_image_PIL(mesh.textures).save(DATA_DIR / "out_apartment.png")

        for i in range(19):
            # random locations in the apartment
            eye = ((np.random.uniform(-6, 0.5), np.random.uniform(-8, 2), 0),)
            at = ((np.random.uniform(-6, 0.5), np.random.uniform(-8, 2), 0),)
            up = ((0, 0, -1),)
            RT = look_at_view_transform(eye=eye, at=at, up=up)
            _render(mesh, f"apartment_eau{i}", RT=RT, use_ambient=True)

        for i in range(12):
            # panning around the inner room from one location
            pan = axis_angle_to_matrix(torch.FloatTensor([0, radians(30 * i), 0]))
            _render(mesh, f"apartment{i}", 1.0, -90, pan, use_ambient=True)

    def test_load_cow(self):
        """
        Load the cow as converted to a single mesh in a glb file.
        """
        glb = DATA_DIR / "cow.glb"
        self.assertTrue(glb.is_file())
        device = torch.device("cuda:0")
        mesh = _load(glb, device=device)
        self.assertEqual(mesh.device, device)

        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.get_bounding_boxes().cpu(), mesh_obj.get_bounding_boxes())

        self.assertClose(
            mesh.textures.verts_uvs_padded().cpu(), mesh_obj.textures.verts_uvs_padded()
        )

        self.assertClose(
            mesh.textures.faces_uvs_padded().cpu(), mesh_obj.textures.faces_uvs_padded()
        )

        self.assertClose(
            mesh.textures.maps_padded().cpu(), mesh_obj.textures.maps_padded()
        )

        if DEBUG:
            texturesuv_image_PIL(mesh.textures).save(DATA_DIR / "out_cow.png")

            image = _render(mesh, "cow", azim=4)
            with Image.open(DATA_DIR / "glb_cow.png") as f:
                expected = np.array(f)

            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.
        """
        glb = DATA_DIR / "cow.glb"
        self.assertTrue(glb.is_file())
        device = torch.device("cuda:0")
        mesh = _load(glb, device=device, include_textures=False)
        self.assertEqual(len(mesh), 1)
        self.assertIsNone(mesh.textures)

        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.get_bounding_boxes().cpu(), mesh_obj.get_bounding_boxes())

        mesh.textures = TexturesVertex(0.5 * torch.ones_like(mesh.verts_padded()))

        image = _render(mesh, "cow_gray")

        with Image.open(DATA_DIR / "glb_cow_gray.png") as f:
            expected = np.array(f)

        self.assertClose(image, expected)

    def test_load_save_load_cow_texturesvertex(self):
        """
        Load the cow as converted to a single mesh in a glb file and then save it to a glb file.
        """

        glb = DATA_DIR / "cow.glb"
        self.assertTrue(glb.is_file())
        device = torch.device("cuda:0")
        mesh = _load(glb, device=device, include_textures=False)
        self.assertEqual(len(mesh), 1)
        self.assertIsNone(mesh.textures)

        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.get_bounding_boxes().cpu(), mesh_obj.get_bounding_boxes())

        mesh.textures = TexturesVertex(0.5 * torch.ones_like(mesh.verts_padded()))

        image = _render(mesh, "cow_gray")

        with Image.open(DATA_DIR / "glb_cow_gray.png") as f:
            expected = np.array(f)

        self.assertClose(image, expected)

        # save the mesh to a glb file
        glb = DATA_DIR / "cow_write_texturesvertex.glb"
        _write(mesh, glb)

        # reload the mesh glb file saved in TexturesVertex format
        glb = DATA_DIR / "cow_write_texturesvertex.glb"
        self.assertTrue(glb.is_file())
        mesh_dash = _load(glb, device=device)
        self.assertEqual(len(mesh_dash), 1)

        self.assertEqual(mesh_dash.faces_packed().shape, (5856, 3))
        self.assertEqual(mesh_dash.verts_packed().shape, (3225, 3))
        self.assertEqual(mesh_dash.textures.verts_features_list()[0].shape, (3225, 3))

        # check the re-rendered image with expected
        image_dash = _render(mesh, "cow_gray_texturesvertex")
        self.assertClose(image_dash, expected)

    def test_save_toy(self):
        """
        Construct a simple mesh and save it to a glb file in TexturesVertex mode.
        """

        example = {}
        example["POSITION"] = torch.tensor(
            [
                [
                    [0.0, 0.0, 0.0],
                    [-1.0, 0.0, 0.0],
                    [-1.0, 0.0, 1.0],
                    [0.0, 0.0, 1.0],
                    [0.0, 1.0, 0.0],
                    [-1.0, 1.0, 0.0],
                    [-1.0, 1.0, 1.0],
                    [0.0, 1.0, 1.0],
                ]
            ]
        )
        example["indices"] = torch.tensor(
            [
                [
                    [1, 4, 2],
                    [4, 3, 2],
                    [3, 7, 2],
                    [7, 6, 2],
                    [3, 4, 7],
                    [4, 8, 7],
                    [8, 5, 7],
                    [5, 6, 7],
                    [5, 2, 6],
                    [5, 1, 2],
                    [1, 5, 4],
                    [5, 8, 4],
                ]
            ]
        )
        example["indices"] -= 1
        example["COLOR_0"] = torch.tensor(
            [
                [
                    [1.0, 0.0, 0.0],
                    [1.0, 0.0, 0.0],
                    [1.0, 0.0, 0.0],
                    [1.0, 0.0, 0.0],
                    [1.0, 0.0, 0.0],
                    [1.0, 0.0, 0.0],
                    [1.0, 0.0, 0.0],
                    [1.0, 0.0, 0.0],
                ]
            ]
        )
        # example['prop'] = {'material':
        #                       {'pbrMetallicRoughness':
        #                           {'baseColorFactor':
        #                                torch.tensor([[0.7, 0.7, 1, 0.5]]),
        #                            'metallicFactor': torch.tensor([1]),
        #                            'roughnessFactor': torch.tensor([0.1])},
        #                    'alphaMode': 'BLEND',
        #                    'doubleSided': True}}

        texture = TexturesVertex(example["COLOR_0"])
        mesh = Meshes(
            verts=example["POSITION"], faces=example["indices"], textures=texture
        )

        glb = DATA_DIR / "example_write_texturesvertex.glb"
        _write(mesh, glb)