mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
b7a316232a
Summary: Make GLB files report their own length correctly. They were off by 28. Reviewed By: davidsonic Differential Revision: D41838340 fbshipit-source-id: 9cd66e8337c142298d5ae1d7c27e51fd812d5c7b
269 lines
8.3 KiB
Python
269 lines
8.3 KiB
Python
# 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)
|