# 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)