mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
26d2cc24c1
Summary: When rendering meshes with Phong shading, interpolate_face_attributes was taking up a nontrivial fraction of the overall shading time. This diff replaces our Python implementation of this function with a CUDA implementation. Reviewed By: nikhilaravi Differential Revision: D21610763 fbshipit-source-id: 2bb362a28f698541812aeab539047264b125ebb8
302 lines
11 KiB
Python
302 lines
11 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
|
|
|
|
import unittest
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from common_testing import TestCaseMixin
|
|
from pytorch3d.renderer.mesh.rasterizer import Fragments
|
|
from pytorch3d.renderer.mesh.texturing import interpolate_texture_map
|
|
from pytorch3d.structures import Meshes, Textures
|
|
from pytorch3d.structures.utils import list_to_padded
|
|
from test_meshes import TestMeshes
|
|
|
|
|
|
class TestTexturing(TestCaseMixin, unittest.TestCase):
|
|
def test_interpolate_texture_map(self):
|
|
barycentric_coords = torch.tensor(
|
|
[[0.5, 0.3, 0.2], [0.3, 0.6, 0.1]], dtype=torch.float32
|
|
).view(1, 1, 1, 2, -1)
|
|
dummy_verts = torch.zeros(4, 3)
|
|
vert_uvs = torch.tensor([[1, 0], [0, 1], [1, 1], [0, 0]], dtype=torch.float32)
|
|
face_uvs = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.int64)
|
|
interpolated_uvs = torch.tensor(
|
|
[[0.5 + 0.2, 0.3 + 0.2], [0.6, 0.3 + 0.6]], dtype=torch.float32
|
|
)
|
|
|
|
# Create a dummy texture map
|
|
H = 2
|
|
W = 2
|
|
x = torch.linspace(0, 1, W).view(1, W).expand(H, W)
|
|
y = torch.linspace(0, 1, H).view(H, 1).expand(H, W)
|
|
tex_map = torch.stack([x, y], dim=2).view(1, H, W, 2)
|
|
pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
|
|
fragments = Fragments(
|
|
pix_to_face=pix_to_face,
|
|
bary_coords=barycentric_coords,
|
|
zbuf=pix_to_face,
|
|
dists=pix_to_face,
|
|
)
|
|
tex = Textures(
|
|
maps=tex_map, faces_uvs=face_uvs[None, ...], verts_uvs=vert_uvs[None, ...]
|
|
)
|
|
meshes = Meshes(verts=[dummy_verts], faces=[face_uvs], textures=tex)
|
|
texels = interpolate_texture_map(fragments, meshes)
|
|
|
|
# Expected output
|
|
pixel_uvs = interpolated_uvs * 2.0 - 1.0
|
|
pixel_uvs = pixel_uvs.view(2, 1, 1, 2)
|
|
tex_map = torch.flip(tex_map, [1])
|
|
tex_map = tex_map.permute(0, 3, 1, 2)
|
|
tex_map = torch.cat([tex_map, tex_map], dim=0)
|
|
expected_out = F.grid_sample(tex_map, pixel_uvs, align_corners=False)
|
|
self.assertTrue(torch.allclose(texels.squeeze(), expected_out.squeeze()))
|
|
|
|
def test_init_rgb_uv_fail(self):
|
|
V = 20
|
|
# Maps has wrong shape
|
|
with self.assertRaisesRegex(ValueError, "maps"):
|
|
Textures(
|
|
maps=torch.ones((5, 16, 16, 3, 4)),
|
|
faces_uvs=torch.randint(size=(5, 10, 3), low=0, high=V),
|
|
verts_uvs=torch.ones((5, V, 2)),
|
|
)
|
|
# faces_uvs has wrong shape
|
|
with self.assertRaisesRegex(ValueError, "faces_uvs"):
|
|
Textures(
|
|
maps=torch.ones((5, 16, 16, 3)),
|
|
faces_uvs=torch.randint(size=(5, 10, 3, 3), low=0, high=V),
|
|
verts_uvs=torch.ones((5, V, 2)),
|
|
)
|
|
# verts_uvs has wrong shape
|
|
with self.assertRaisesRegex(ValueError, "verts_uvs"):
|
|
Textures(
|
|
maps=torch.ones((5, 16, 16, 3)),
|
|
faces_uvs=torch.randint(size=(5, 10, 3), low=0, high=V),
|
|
verts_uvs=torch.ones((5, V, 2, 3)),
|
|
)
|
|
# verts_rgb has wrong shape
|
|
with self.assertRaisesRegex(ValueError, "verts_rgb"):
|
|
Textures(verts_rgb=torch.ones((5, 16, 16, 3)))
|
|
|
|
# maps provided without verts/faces uvs
|
|
with self.assertRaisesRegex(ValueError, "faces_uvs and verts_uvs are required"):
|
|
Textures(maps=torch.ones((5, 16, 16, 3)))
|
|
|
|
def test_padded_to_packed(self):
|
|
N = 2
|
|
# Case where each face in the mesh has 3 unique uv vertex indices
|
|
# - i.e. even if a vertex is shared between multiple faces it will
|
|
# have a unique uv coordinate for each face.
|
|
faces_uvs_list = [
|
|
torch.tensor([[0, 1, 2], [3, 5, 4], [7, 6, 8]]),
|
|
torch.tensor([[0, 1, 2], [3, 4, 5]]),
|
|
] # (N, 3, 3)
|
|
verts_uvs_list = [torch.ones(9, 2), torch.ones(6, 2)]
|
|
faces_uvs_padded = list_to_padded(faces_uvs_list, pad_value=-1)
|
|
verts_uvs_padded = list_to_padded(verts_uvs_list)
|
|
tex = Textures(
|
|
maps=torch.ones((N, 16, 16, 3)),
|
|
faces_uvs=faces_uvs_padded,
|
|
verts_uvs=verts_uvs_padded,
|
|
)
|
|
|
|
# This is set inside Meshes when textures is passed as an input.
|
|
# Here we set _num_faces_per_mesh and _num_verts_per_mesh explicity.
|
|
tex1 = tex.clone()
|
|
tex1._num_faces_per_mesh = faces_uvs_padded.gt(-1).all(-1).sum(-1).tolist()
|
|
tex1._num_verts_per_mesh = torch.tensor([5, 4])
|
|
faces_packed = tex1.faces_uvs_packed()
|
|
verts_packed = tex1.verts_uvs_packed()
|
|
faces_list = tex1.faces_uvs_list()
|
|
verts_list = tex1.verts_uvs_list()
|
|
|
|
for f1, f2 in zip(faces_uvs_list, faces_list):
|
|
self.assertTrue((f1 == f2).all().item())
|
|
|
|
for f, v1, v2 in zip(faces_list, verts_list, verts_uvs_list):
|
|
idx = f.unique()
|
|
self.assertTrue((v1[idx] == v2).all().item())
|
|
|
|
self.assertTrue(faces_packed.shape == (3 + 2, 3))
|
|
|
|
# verts_packed is just flattened verts_padded.
|
|
# split sizes are not used for verts_uvs.
|
|
self.assertTrue(verts_packed.shape == (9 * 2, 2))
|
|
|
|
# Case where num_faces_per_mesh is not set
|
|
tex2 = tex.clone()
|
|
faces_packed = tex2.faces_uvs_packed()
|
|
verts_packed = tex2.verts_uvs_packed()
|
|
faces_list = tex2.faces_uvs_list()
|
|
verts_list = tex2.verts_uvs_list()
|
|
|
|
# Packed is just flattened padded as num_faces_per_mesh
|
|
# has not been provided.
|
|
self.assertTrue(verts_packed.shape == (9 * 2, 2))
|
|
self.assertTrue(faces_packed.shape == (3 * 2, 3))
|
|
|
|
for i in range(N):
|
|
self.assertTrue(
|
|
(faces_list[i] == faces_uvs_padded[i, ...].squeeze()).all().item()
|
|
)
|
|
|
|
for i in range(N):
|
|
self.assertTrue(
|
|
(verts_list[i] == verts_uvs_padded[i, ...].squeeze()).all().item()
|
|
)
|
|
|
|
def test_clone(self):
|
|
V = 20
|
|
tex = Textures(
|
|
maps=torch.ones((5, 16, 16, 3)),
|
|
faces_uvs=torch.randint(size=(5, 10, 3), low=0, high=V),
|
|
verts_uvs=torch.ones((5, V, 2)),
|
|
)
|
|
tex_cloned = tex.clone()
|
|
self.assertSeparate(tex._faces_uvs_padded, tex_cloned._faces_uvs_padded)
|
|
self.assertSeparate(tex._verts_uvs_padded, tex_cloned._verts_uvs_padded)
|
|
self.assertSeparate(tex._maps_padded, tex_cloned._maps_padded)
|
|
|
|
def test_getitem(self):
|
|
N = 5
|
|
V = 20
|
|
source = {
|
|
"maps": torch.rand(size=(N, 16, 16, 3)),
|
|
"faces_uvs": torch.randint(size=(N, 10, 3), low=0, high=V),
|
|
"verts_uvs": torch.rand((N, V, 2)),
|
|
}
|
|
tex = Textures(
|
|
maps=source["maps"],
|
|
faces_uvs=source["faces_uvs"],
|
|
verts_uvs=source["verts_uvs"],
|
|
)
|
|
|
|
verts = torch.rand(size=(N, V, 3))
|
|
faces = torch.randint(size=(N, 10, 3), high=V)
|
|
meshes = Meshes(verts=verts, faces=faces, textures=tex)
|
|
|
|
def tryindex(index):
|
|
tex2 = tex[index]
|
|
meshes2 = meshes[index]
|
|
tex_from_meshes = meshes2.textures
|
|
for item in source:
|
|
basic = source[item][index]
|
|
from_texture = getattr(tex2, item + "_padded")()
|
|
from_meshes = getattr(tex_from_meshes, item + "_padded")()
|
|
if isinstance(index, int):
|
|
basic = basic[None]
|
|
self.assertClose(basic, from_texture)
|
|
self.assertClose(basic, from_meshes)
|
|
self.assertEqual(
|
|
from_texture.ndim, getattr(tex, item + "_padded")().ndim
|
|
)
|
|
if item == "faces_uvs":
|
|
faces_uvs_list = tex_from_meshes.faces_uvs_list()
|
|
self.assertEqual(basic.shape[0], len(faces_uvs_list))
|
|
for i, faces_uvs in enumerate(faces_uvs_list):
|
|
self.assertClose(faces_uvs, basic[i])
|
|
|
|
tryindex(2)
|
|
tryindex(slice(0, 2, 1))
|
|
index = torch.tensor([1, 0, 1, 0, 0], dtype=torch.bool)
|
|
tryindex(index)
|
|
index = torch.tensor([0, 0, 0, 0, 0], dtype=torch.bool)
|
|
tryindex(index)
|
|
index = torch.tensor([1, 2], dtype=torch.int64)
|
|
tryindex(index)
|
|
tryindex([2, 4])
|
|
|
|
def test_to(self):
|
|
V = 20
|
|
tex = Textures(
|
|
maps=torch.ones((5, 16, 16, 3)),
|
|
faces_uvs=torch.randint(size=(5, 10, 3), low=0, high=V),
|
|
verts_uvs=torch.ones((5, V, 2)),
|
|
)
|
|
device = torch.device("cuda:0")
|
|
tex = tex.to(device)
|
|
self.assertTrue(tex._faces_uvs_padded.device == device)
|
|
self.assertTrue(tex._verts_uvs_padded.device == device)
|
|
self.assertTrue(tex._maps_padded.device == device)
|
|
|
|
def test_extend(self):
|
|
B = 10
|
|
mesh = TestMeshes.init_mesh(B, 30, 50)
|
|
V = mesh._V
|
|
F = mesh._F
|
|
|
|
# 1. Texture uvs
|
|
tex_uv = Textures(
|
|
maps=torch.randn((B, 16, 16, 3)),
|
|
faces_uvs=torch.randint(size=(B, F, 3), low=0, high=V),
|
|
verts_uvs=torch.randn((B, V, 2)),
|
|
)
|
|
tex_mesh = Meshes(
|
|
verts=mesh.verts_padded(), faces=mesh.faces_padded(), textures=tex_uv
|
|
)
|
|
N = 20
|
|
new_mesh = tex_mesh.extend(N)
|
|
|
|
self.assertEqual(len(tex_mesh) * N, len(new_mesh))
|
|
|
|
tex_init = tex_mesh.textures
|
|
new_tex = new_mesh.textures
|
|
|
|
for i in range(len(tex_mesh)):
|
|
for n in range(N):
|
|
self.assertClose(
|
|
tex_init.faces_uvs_list()[i], new_tex.faces_uvs_list()[i * N + n]
|
|
)
|
|
self.assertClose(
|
|
tex_init.verts_uvs_list()[i], new_tex.verts_uvs_list()[i * N + n]
|
|
)
|
|
self.assertAllSeparate(
|
|
[
|
|
tex_init.faces_uvs_padded(),
|
|
new_tex.faces_uvs_padded(),
|
|
tex_init.verts_uvs_padded(),
|
|
new_tex.verts_uvs_padded(),
|
|
tex_init.maps_padded(),
|
|
new_tex.maps_padded(),
|
|
]
|
|
)
|
|
|
|
self.assertIsNone(new_tex.verts_rgb_list())
|
|
self.assertIsNone(new_tex.verts_rgb_padded())
|
|
self.assertIsNone(new_tex.verts_rgb_packed())
|
|
|
|
# 2. Texture vertex RGB
|
|
tex_rgb = Textures(verts_rgb=torch.randn((B, V, 3)))
|
|
tex_mesh_rgb = Meshes(
|
|
verts=mesh.verts_padded(), faces=mesh.faces_padded(), textures=tex_rgb
|
|
)
|
|
N = 20
|
|
new_mesh_rgb = tex_mesh_rgb.extend(N)
|
|
|
|
self.assertEqual(len(tex_mesh_rgb) * N, len(new_mesh_rgb))
|
|
|
|
tex_init = tex_mesh_rgb.textures
|
|
new_tex = new_mesh_rgb.textures
|
|
|
|
for i in range(len(tex_mesh_rgb)):
|
|
for n in range(N):
|
|
self.assertClose(
|
|
tex_init.verts_rgb_list()[i], new_tex.verts_rgb_list()[i * N + n]
|
|
)
|
|
self.assertAllSeparate(
|
|
[tex_init.verts_rgb_padded(), new_tex.verts_rgb_padded()]
|
|
)
|
|
|
|
self.assertIsNone(new_tex.verts_uvs_padded())
|
|
self.assertIsNone(new_tex.verts_uvs_list())
|
|
self.assertIsNone(new_tex.verts_uvs_packed())
|
|
self.assertIsNone(new_tex.faces_uvs_padded())
|
|
self.assertIsNone(new_tex.faces_uvs_list())
|
|
self.assertIsNone(new_tex.faces_uvs_packed())
|
|
|
|
# 3. Error
|
|
with self.assertRaises(ValueError):
|
|
tex_mesh.extend(N=-1)
|