getitem for textures

Summary: Make Meshes.__getitem__ carry texture information to the new mesh.

Reviewed By: gkioxari

Differential Revision: D20283976

fbshipit-source-id: d9ee0580c11ac5b4384df9d8158a07e6eb8d00fe
This commit is contained in:
Jeremy Reizenstein
2020-03-11 07:42:35 -07:00
committed by Facebook GitHub Bot
parent 5a1d7143d8
commit fb97ab104e
3 changed files with 71 additions and 2 deletions

View File

@@ -167,6 +167,55 @@ class TestTexturing(TestCaseMixin, unittest.TestCase):
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(