test_texturing.py

# 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.textures import (
    TexturesAtlas,
    TexturesUV,
    TexturesVertex,
    _list_to_padded_wrapper,
    pack_rectangles,
)
from pytorch3d.structures import Meshes, list_to_packed, packed_to_list
from test_meshes import init_mesh


def tryindex(self, index, tex, meshes, source):
    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]

        if len(basic) == 0:
            self.assertEqual(len(from_texture), 0)
            self.assertEqual(len(from_meshes), 0)
        else:
            self.assertClose(basic, from_texture)
            self.assertClose(basic, from_meshes)
            self.assertEqual(from_texture.ndim, getattr(tex, item + "_padded")().ndim)
            item_list = getattr(tex_from_meshes, item + "_list")()
            self.assertEqual(basic.shape[0], len(item_list))
            for i, elem in enumerate(item_list):
                self.assertClose(elem, basic[i])


class TestTexturesVertex(TestCaseMixin, unittest.TestCase):
    def test_sample_vertex_textures(self):
        """
        This tests both interpolate_vertex_colors as well as
        interpolate_face_attributes.
        """
        verts = torch.randn((4, 3), dtype=torch.float32)
        faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
        vert_tex = torch.tensor(
            [[0, 1, 0], [0, 1, 1], [1, 1, 0], [1, 1, 1]], dtype=torch.float32
        )
        verts_features = vert_tex
        tex = TexturesVertex(verts_features=[verts_features])
        mesh = Meshes(verts=[verts], faces=[faces], textures=tex)
        pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
        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)
        expected_vals = torch.tensor(
            [[0.5, 1.0, 0.3], [0.3, 1.0, 0.9]], dtype=torch.float32
        ).view(1, 1, 1, 2, -1)
        fragments = Fragments(
            pix_to_face=pix_to_face,
            bary_coords=barycentric_coords,
            zbuf=torch.ones_like(pix_to_face),
            dists=torch.ones_like(pix_to_face),
        )
        # sample_textures calls interpolate_vertex_colors
        texels = mesh.sample_textures(fragments)
        self.assertTrue(torch.allclose(texels, expected_vals[None, :]))

    def test_sample_vertex_textures_grad(self):
        verts = torch.randn((4, 3), dtype=torch.float32)
        faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
        vert_tex = torch.tensor(
            [[0, 1, 0], [0, 1, 1], [1, 1, 0], [1, 1, 1]],
            dtype=torch.float32,
            requires_grad=True,
        )
        verts_features = vert_tex
        tex = TexturesVertex(verts_features=[verts_features])
        mesh = Meshes(verts=[verts], faces=[faces], textures=tex)
        pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
        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)
        fragments = Fragments(
            pix_to_face=pix_to_face,
            bary_coords=barycentric_coords,
            zbuf=torch.ones_like(pix_to_face),
            dists=torch.ones_like(pix_to_face),
        )
        grad_vert_tex = torch.tensor(
            [[0.3, 0.3, 0.3], [0.9, 0.9, 0.9], [0.5, 0.5, 0.5], [0.3, 0.3, 0.3]],
            dtype=torch.float32,
        )
        texels = mesh.sample_textures(fragments)
        texels.sum().backward()
        self.assertTrue(hasattr(vert_tex, "grad"))
        self.assertTrue(torch.allclose(vert_tex.grad, grad_vert_tex[None, :]))

    def test_textures_vertex_init_fail(self):
        # Incorrect sized tensors
        with self.assertRaisesRegex(ValueError, "verts_features"):
            TexturesVertex(verts_features=torch.rand(size=(5, 10)))

        # Not a list or a tensor
        with self.assertRaisesRegex(ValueError, "verts_features"):
            TexturesVertex(verts_features=(1, 1, 1))

    def test_faces_verts_textures(self):
        device = torch.device("cuda:0")
        verts = torch.randn((2, 4, 3), dtype=torch.float32, device=device)
        faces = torch.tensor(
            [[[2, 1, 0], [3, 1, 0]], [[1, 3, 0], [2, 1, 3]]],
            dtype=torch.int64,
            device=device,
        )

        # define TexturesVertex
        verts_texture = torch.rand(verts.shape)
        textures = TexturesVertex(verts_features=verts_texture)

        # compute packed faces
        ff = faces.unbind(0)
        faces_packed = torch.cat([ff[0], ff[1] + verts.shape[1]])

        # face verts textures
        faces_verts_texts = textures.faces_verts_textures_packed(faces_packed)

        verts_texts_packed = torch.cat(verts_texture.unbind(0))
        faces_verts_texts_packed = verts_texts_packed[faces_packed]

        self.assertClose(faces_verts_texts_packed, faces_verts_texts)

    def test_clone(self):
        tex = TexturesVertex(verts_features=torch.rand(size=(10, 100, 128)))
        tex.verts_features_list()
        tex_cloned = tex.clone()
        self.assertSeparate(
            tex._verts_features_padded, tex_cloned._verts_features_padded
        )
        self.assertClose(tex._verts_features_padded, tex_cloned._verts_features_padded)
        self.assertSeparate(tex.valid, tex_cloned.valid)
        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
        for i in range(tex._N):
            self.assertSeparate(
                tex._verts_features_list[i], tex_cloned._verts_features_list[i]
            )
            self.assertClose(
                tex._verts_features_list[i], tex_cloned._verts_features_list[i]
            )

    def test_detach(self):
        tex = TexturesVertex(
            verts_features=torch.rand(size=(10, 100, 128), requires_grad=True)
        )
        tex.verts_features_list()
        tex_detached = tex.detach()
        self.assertFalse(tex_detached._verts_features_padded.requires_grad)
        self.assertClose(
            tex_detached._verts_features_padded, tex._verts_features_padded
        )
        for i in range(tex._N):
            self.assertClose(
                tex._verts_features_list[i], tex_detached._verts_features_list[i]
            )
            self.assertFalse(tex_detached._verts_features_list[i].requires_grad)

    def test_extend(self):
        B = 10
        mesh = init_mesh(B, 30, 50)
        V = mesh._V
        tex_uv = TexturesVertex(verts_features=torch.randn((B, V, 3)))
        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.verts_features_list()[i],
                    new_tex.verts_features_list()[i * N + n],
                )
                self.assertClose(
                    tex_init._num_faces_per_mesh[i],
                    new_tex._num_faces_per_mesh[i * N + n],
                )

        self.assertAllSeparate(
            [tex_init.verts_features_padded(), new_tex.verts_features_padded()]
        )

        with self.assertRaises(ValueError):
            tex_mesh.extend(N=-1)

    def test_padded_to_packed(self):
        # 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.
        num_verts_per_mesh = [9, 6]
        D = 10
        verts_features_list = [torch.rand(v, D) for v in num_verts_per_mesh]
        verts_features_packed = list_to_packed(verts_features_list)[0]
        verts_features_list = packed_to_list(verts_features_packed, num_verts_per_mesh)
        tex = TexturesVertex(verts_features=verts_features_list)

        # 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_verts_per_mesh = num_verts_per_mesh
        verts_packed = tex1.verts_features_packed()
        verts_verts_list = tex1.verts_features_list()
        verts_padded = tex1.verts_features_padded()

        for f1, f2 in zip(verts_verts_list, verts_features_list):
            self.assertTrue((f1 == f2).all().item())

        self.assertTrue(verts_packed.shape == (sum(num_verts_per_mesh), D))
        self.assertTrue(verts_padded.shape == (2, 9, D))

        # Case where num_verts_per_mesh is not set and textures
        # are initialized with a padded tensor.
        tex2 = TexturesVertex(verts_features=verts_padded)
        verts_packed = tex2.verts_features_packed()
        verts_list = tex2.verts_features_list()

        # Packed is just flattened padded as num_verts_per_mesh
        # has not been provided.
        self.assertTrue(verts_packed.shape == (9 * 2, D))

        for i, (f1, f2) in enumerate(zip(verts_list, verts_features_list)):
            n = num_verts_per_mesh[i]
            self.assertTrue((f1[:n] == f2).all().item())

    def test_getitem(self):
        N = 5
        V = 20
        source = {"verts_features": torch.randn(size=(N, 10, 128))}
        tex = TexturesVertex(verts_features=source["verts_features"])

        verts = torch.rand(size=(N, V, 3))
        faces = torch.randint(size=(N, 10, 3), high=V)
        meshes = Meshes(verts=verts, faces=faces, textures=tex)

        tryindex(self, 2, tex, meshes, source)
        tryindex(self, slice(0, 2, 1), tex, meshes, source)
        index = torch.tensor([1, 0, 1, 0, 0], dtype=torch.bool)
        tryindex(self, index, tex, meshes, source)
        index = torch.tensor([0, 0, 0, 0, 0], dtype=torch.bool)
        tryindex(self, index, tex, meshes, source)
        index = torch.tensor([1, 2], dtype=torch.int64)
        tryindex(self, index, tex, meshes, source)
        tryindex(self, [2, 4], tex, meshes, source)


class TestTexturesAtlas(TestCaseMixin, unittest.TestCase):
    def test_sample_texture_atlas(self):
        N, F, R = 1, 2, 2
        verts = torch.randn((4, 3), dtype=torch.float32)
        faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
        faces_atlas = torch.rand(size=(N, F, R, R, 3))
        tex = TexturesAtlas(atlas=faces_atlas)
        mesh = Meshes(verts=[verts], faces=[faces], textures=tex)
        pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
        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)
        expected_vals = torch.tensor(
            [[0.5, 1.0, 0.3], [0.3, 1.0, 0.9]], dtype=torch.float32
        )
        expected_vals = torch.zeros((1, 1, 1, 2, 3), dtype=torch.float32)
        expected_vals[..., 0, :] = faces_atlas[0, 0, 0, 1, ...]
        expected_vals[..., 1, :] = faces_atlas[0, 1, 1, 0, ...]

        fragments = Fragments(
            pix_to_face=pix_to_face,
            bary_coords=barycentric_coords,
            zbuf=torch.ones_like(pix_to_face),
            dists=torch.ones_like(pix_to_face),
        )
        texels = mesh.textures.sample_textures(fragments)
        self.assertTrue(torch.allclose(texels, expected_vals))

    def test_textures_atlas_grad(self):
        N, F, R = 1, 2, 2
        verts = torch.randn((4, 3), dtype=torch.float32)
        faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
        faces_atlas = torch.rand(size=(N, F, R, R, 3), requires_grad=True)
        tex = TexturesAtlas(atlas=faces_atlas)
        mesh = Meshes(verts=[verts], faces=[faces], textures=tex)
        pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
        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)
        fragments = Fragments(
            pix_to_face=pix_to_face,
            bary_coords=barycentric_coords,
            zbuf=torch.ones_like(pix_to_face),
            dists=torch.ones_like(pix_to_face),
        )
        texels = mesh.textures.sample_textures(fragments)
        grad_tex = torch.rand_like(texels)
        grad_expected = torch.zeros_like(faces_atlas)
        grad_expected[0, 0, 0, 1, :] = grad_tex[..., 0:1, :]
        grad_expected[0, 1, 1, 0, :] = grad_tex[..., 1:2, :]
        texels.backward(grad_tex)
        self.assertTrue(hasattr(faces_atlas, "grad"))
        self.assertTrue(torch.allclose(faces_atlas.grad, grad_expected))

    def test_textures_atlas_init_fail(self):
        # Incorrect sized tensors
        with self.assertRaisesRegex(ValueError, "atlas"):
            TexturesAtlas(atlas=torch.rand(size=(5, 10, 3)))

        # Not a list or a tensor
        with self.assertRaisesRegex(ValueError, "atlas"):
            TexturesAtlas(atlas=(1, 1, 1))

    def test_faces_verts_textures(self):
        device = torch.device("cuda:0")
        N, F, R = 2, 2, 8
        num_faces = torch.randint(low=1, high=F, size=(N,))
        faces_atlas = [
            torch.rand(size=(num_faces[i].item(), R, R, 3), device=device)
            for i in range(N)
        ]
        tex = TexturesAtlas(atlas=faces_atlas)

        # faces_verts naive
        faces_verts = []
        for n in range(N):
            ff = num_faces[n].item()
            temp = torch.zeros(ff, 3, 3)
            for f in range(ff):
                t0 = faces_atlas[n][f, 0, -1]  # for v0, bary = (1, 0)
                t1 = faces_atlas[n][f, -1, 0]  # for v1, bary = (0, 1)
                t2 = faces_atlas[n][f, 0, 0]  # for v2, bary = (0, 0)
                temp[f, 0] = t0
                temp[f, 1] = t1
                temp[f, 2] = t2
            faces_verts.append(temp)
        faces_verts = torch.cat(faces_verts, 0)

        self.assertClose(faces_verts, tex.faces_verts_textures_packed().cpu())

    def test_clone(self):
        tex = TexturesAtlas(atlas=torch.rand(size=(1, 10, 2, 2, 3)))
        tex.atlas_list()
        tex_cloned = tex.clone()
        self.assertSeparate(tex._atlas_padded, tex_cloned._atlas_padded)
        self.assertClose(tex._atlas_padded, tex_cloned._atlas_padded)
        self.assertSeparate(tex.valid, tex_cloned.valid)
        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
        for i in range(tex._N):
            self.assertSeparate(tex._atlas_list[i], tex_cloned._atlas_list[i])
            self.assertClose(tex._atlas_list[i], tex_cloned._atlas_list[i])

    def test_detach(self):
        tex = TexturesAtlas(atlas=torch.rand(size=(1, 10, 2, 2, 3), requires_grad=True))
        tex.atlas_list()
        tex_detached = tex.detach()
        self.assertFalse(tex_detached._atlas_padded.requires_grad)
        self.assertClose(tex_detached._atlas_padded, tex._atlas_padded)
        for i in range(tex._N):
            self.assertFalse(tex_detached._atlas_list[i].requires_grad)
            self.assertClose(tex._atlas_list[i], tex_detached._atlas_list[i])

    def test_extend(self):
        B = 10
        mesh = init_mesh(B, 30, 50)
        F = mesh._F
        tex_uv = TexturesAtlas(atlas=torch.randn((B, F, 2, 2, 3)))
        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.atlas_list()[i], new_tex.atlas_list()[i * N + n]
                )
                self.assertClose(
                    tex_init._num_faces_per_mesh[i],
                    new_tex._num_faces_per_mesh[i * N + n],
                )

        self.assertAllSeparate([tex_init.atlas_padded(), new_tex.atlas_padded()])

        with self.assertRaises(ValueError):
            tex_mesh.extend(N=-1)

    def test_padded_to_packed(self):
        # 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.
        R = 2
        N = 20
        num_faces_per_mesh = torch.randint(size=(N,), low=0, high=30)
        atlas_list = [torch.rand(f, R, R, 3) for f in num_faces_per_mesh]
        tex = TexturesAtlas(atlas=atlas_list)

        # This is set inside Meshes when textures is passed as an input.
        # Here we set _num_faces_per_mesh explicity.
        tex1 = tex.clone()
        tex1._num_faces_per_mesh = num_faces_per_mesh.tolist()
        atlas_packed = tex1.atlas_packed()
        atlas_list_new = tex1.atlas_list()
        atlas_padded = tex1.atlas_padded()

        for f1, f2 in zip(atlas_list_new, atlas_list):
            self.assertTrue((f1 == f2).all().item())

        sum_F = num_faces_per_mesh.sum()
        max_F = num_faces_per_mesh.max().item()
        self.assertTrue(atlas_packed.shape == (sum_F, R, R, 3))
        self.assertTrue(atlas_padded.shape == (N, max_F, R, R, 3))

        # Case where num_faces_per_mesh is not set and textures
        # are initialized with a padded tensor.
        atlas_list_padded = _list_to_padded_wrapper(atlas_list)
        tex2 = TexturesAtlas(atlas=atlas_list_padded)
        atlas_packed = tex2.atlas_packed()
        atlas_list_new = tex2.atlas_list()

        # Packed is just flattened padded as num_faces_per_mesh
        # has not been provided.
        self.assertTrue(atlas_packed.shape == (N * max_F, R, R, 3))

        for i, (f1, f2) in enumerate(zip(atlas_list_new, atlas_list)):
            n = num_faces_per_mesh[i]
            self.assertTrue((f1[:n] == f2).all().item())

    def test_getitem(self):
        N = 5
        V = 20
        source = {"atlas": torch.randn(size=(N, 10, 4, 4, 3))}
        tex = TexturesAtlas(atlas=source["atlas"])

        verts = torch.rand(size=(N, V, 3))
        faces = torch.randint(size=(N, 10, 3), high=V)
        meshes = Meshes(verts=verts, faces=faces, textures=tex)

        tryindex(self, 2, tex, meshes, source)
        tryindex(self, slice(0, 2, 1), tex, meshes, source)
        index = torch.tensor([1, 0, 1, 0, 0], dtype=torch.bool)
        tryindex(self, index, tex, meshes, source)
        index = torch.tensor([0, 0, 0, 0, 0], dtype=torch.bool)
        tryindex(self, index, tex, meshes, source)
        index = torch.tensor([1, 2], dtype=torch.int64)
        tryindex(self, index, tex, meshes, source)
        tryindex(self, [2, 4], tex, meshes, source)


class TestTexturesUV(TestCaseMixin, unittest.TestCase):
    def setUp(self) -> None:
        super().setUp()
        torch.manual_seed(42)

    def test_sample_textures_uv(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,
        )

        for align_corners in [True, False]:
            tex = TexturesUV(
                maps=tex_map,
                faces_uvs=[face_uvs],
                verts_uvs=[vert_uvs],
                align_corners=align_corners,
            )
            meshes = Meshes(verts=[dummy_verts], faces=[face_uvs], textures=tex)
            mesh_textures = meshes.textures
            texels = mesh_textures.sample_textures(fragments)

            # 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]).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=align_corners, padding_mode="border"
            )
            self.assertTrue(torch.allclose(texels.squeeze(), expected_out.squeeze()))

    def test_textures_uv_init_fail(self):
        # Maps has wrong shape
        with self.assertRaisesRegex(ValueError, "maps"):
            TexturesUV(
                maps=torch.ones((5, 16, 16, 3, 4)),
                faces_uvs=torch.rand(size=(5, 10, 3)),
                verts_uvs=torch.rand(size=(5, 15, 2)),
            )

        # faces_uvs has wrong shape
        with self.assertRaisesRegex(ValueError, "faces_uvs"):
            TexturesUV(
                maps=torch.ones((5, 16, 16, 3)),
                faces_uvs=torch.rand(size=(5, 10, 3, 3)),
                verts_uvs=torch.rand(size=(5, 15, 2)),
            )

        # verts_uvs has wrong shape
        with self.assertRaisesRegex(ValueError, "verts_uvs"):
            TexturesUV(
                maps=torch.ones((5, 16, 16, 3)),
                faces_uvs=torch.rand(size=(5, 10, 3)),
                verts_uvs=torch.rand(size=(5, 15, 2, 3)),
            )

        # verts has different batch dim to faces
        with self.assertRaisesRegex(ValueError, "verts_uvs"):
            TexturesUV(
                maps=torch.ones((5, 16, 16, 3)),
                faces_uvs=torch.rand(size=(5, 10, 3)),
                verts_uvs=torch.rand(size=(8, 15, 2)),
            )

        # maps has different batch dim to faces
        with self.assertRaisesRegex(ValueError, "maps"):
            TexturesUV(
                maps=torch.ones((8, 16, 16, 3)),
                faces_uvs=torch.rand(size=(5, 10, 3)),
                verts_uvs=torch.rand(size=(5, 15, 2)),
            )

        # verts on different device to faces
        with self.assertRaisesRegex(ValueError, "verts_uvs"):
            TexturesUV(
                maps=torch.ones((5, 16, 16, 3)),
                faces_uvs=torch.rand(size=(5, 10, 3)),
                verts_uvs=torch.rand(size=(5, 15, 2, 3), device="cuda"),
            )

        # maps on different device to faces
        with self.assertRaisesRegex(ValueError, "map"):
            TexturesUV(
                maps=torch.ones((5, 16, 16, 3), device="cuda"),
                faces_uvs=torch.rand(size=(5, 10, 3)),
                verts_uvs=torch.rand(size=(5, 15, 2)),
            )

    def test_faces_verts_textures(self):
        device = torch.device("cuda:0")
        N, V, F, H, W = 2, 5, 12, 8, 8
        vert_uvs = torch.rand((N, V, 2), dtype=torch.float32, device=device)
        face_uvs = torch.randint(
            high=V, size=(N, F, 3), dtype=torch.int64, device=device
        )
        maps = torch.rand((N, H, W, 3), dtype=torch.float32, device=device)

        tex = TexturesUV(maps=maps, verts_uvs=vert_uvs, faces_uvs=face_uvs)

        # naive faces_verts_textures
        faces_verts_texs = []
        for n in range(N):
            temp = torch.zeros((F, 3, 3), device=device, dtype=torch.float32)
            for f in range(F):
                uv0 = vert_uvs[n, face_uvs[n, f, 0]]
                uv1 = vert_uvs[n, face_uvs[n, f, 1]]
                uv2 = vert_uvs[n, face_uvs[n, f, 2]]

                idx = torch.stack((uv0, uv1, uv2), dim=0).view(1, 1, 3, 2)  # 1x1x3x2
                idx = idx * 2.0 - 1.0
                imap = maps[n].view(1, H, W, 3).permute(0, 3, 1, 2)  # 1x3xHxW
                imap = torch.flip(imap, [2])

                texts = torch.nn.functional.grid_sample(
                    imap,
                    idx,
                    align_corners=tex.align_corners,
                    padding_mode=tex.padding_mode,
                )  # 1x3x1x3
                temp[f] = texts[0, :, 0, :].permute(1, 0)
            faces_verts_texs.append(temp)
        faces_verts_texs = torch.cat(faces_verts_texs, 0)

        self.assertClose(faces_verts_texs, tex.faces_verts_textures_packed())

    def test_clone(self):
        tex = TexturesUV(
            maps=torch.ones((5, 16, 16, 3)),
            faces_uvs=torch.rand(size=(5, 10, 3)),
            verts_uvs=torch.rand(size=(5, 15, 2)),
        )
        tex.faces_uvs_list()
        tex.verts_uvs_list()
        tex_cloned = tex.clone()
        self.assertSeparate(tex._faces_uvs_padded, tex_cloned._faces_uvs_padded)
        self.assertClose(tex._faces_uvs_padded, tex_cloned._faces_uvs_padded)
        self.assertSeparate(tex._verts_uvs_padded, tex_cloned._verts_uvs_padded)
        self.assertClose(tex._verts_uvs_padded, tex_cloned._verts_uvs_padded)
        self.assertSeparate(tex._maps_padded, tex_cloned._maps_padded)
        self.assertClose(tex._maps_padded, tex_cloned._maps_padded)
        self.assertSeparate(tex.valid, tex_cloned.valid)
        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
        for i in range(tex._N):
            self.assertSeparate(tex._faces_uvs_list[i], tex_cloned._faces_uvs_list[i])
            self.assertClose(tex._faces_uvs_list[i], tex_cloned._faces_uvs_list[i])
            self.assertSeparate(tex._verts_uvs_list[i], tex_cloned._verts_uvs_list[i])
            self.assertClose(tex._verts_uvs_list[i], tex_cloned._verts_uvs_list[i])
            # tex._maps_list is not use anywhere so it's not stored. We call it explicitly
            self.assertSeparate(tex.maps_list()[i], tex_cloned.maps_list()[i])
            self.assertClose(tex.maps_list()[i], tex_cloned.maps_list()[i])

    def test_detach(self):
        tex = TexturesUV(
            maps=torch.ones((5, 16, 16, 3), requires_grad=True),
            faces_uvs=torch.rand(size=(5, 10, 3)),
            verts_uvs=torch.rand(size=(5, 15, 2)),
        )
        tex.faces_uvs_list()
        tex.verts_uvs_list()
        tex_detached = tex.detach()
        self.assertFalse(tex_detached._maps_padded.requires_grad)
        self.assertClose(tex._maps_padded, tex_detached._maps_padded)
        self.assertFalse(tex_detached._verts_uvs_padded.requires_grad)
        self.assertClose(tex._verts_uvs_padded, tex_detached._verts_uvs_padded)
        self.assertFalse(tex_detached._faces_uvs_padded.requires_grad)
        self.assertClose(tex._faces_uvs_padded, tex_detached._faces_uvs_padded)
        for i in range(tex._N):
            self.assertFalse(tex_detached._verts_uvs_list[i].requires_grad)
            self.assertClose(tex._verts_uvs_list[i], tex_detached._verts_uvs_list[i])
            self.assertFalse(tex_detached._faces_uvs_list[i].requires_grad)
            self.assertClose(tex._faces_uvs_list[i], tex_detached._faces_uvs_list[i])
            # tex._maps_list is not use anywhere so it's not stored. We call it explicitly
            self.assertFalse(tex_detached.maps_list()[i].requires_grad)
            self.assertClose(tex.maps_list()[i], tex_detached.maps_list()[i])

    def test_extend(self):
        B = 5
        mesh = init_mesh(B, 30, 50)
        V = mesh._V
        num_faces = mesh.num_faces_per_mesh()
        num_verts = mesh.num_verts_per_mesh()
        faces_uvs_list = [torch.randint(size=(f, 3), low=0, high=V) for f in num_faces]
        verts_uvs_list = [torch.rand(v, 2) for v in num_verts]
        tex_uv = TexturesUV(
            maps=torch.ones((B, 16, 16, 3)),
            faces_uvs=faces_uvs_list,
            verts_uvs=verts_uvs_list,
        )
        tex_mesh = Meshes(
            verts=mesh.verts_list(), faces=mesh.faces_list(), textures=tex_uv
        )
        N = 2
        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

        new_tex_num_verts = new_mesh.num_verts_per_mesh()
        for i in range(len(tex_mesh)):
            for n in range(N):
                tex_nv = new_tex_num_verts[i * N + n]
                self.assertClose(
                    # The original textures were initialized using
                    # verts uvs list
                    tex_init.verts_uvs_list()[i],
                    # In the new textures, the verts_uvs are initialized
                    # from padded. The verts per mesh are not used to
                    # convert from padded to list. See TexturesUV for an
                    # explanation.
                    new_tex.verts_uvs_list()[i * N + n][:tex_nv, ...],
                )
                self.assertClose(
                    tex_init.faces_uvs_list()[i], new_tex.faces_uvs_list()[i * N + n]
                )
                self.assertClose(
                    tex_init.maps_padded()[i, ...], new_tex.maps_padded()[i * N + n]
                )
                self.assertClose(
                    tex_init._num_faces_per_mesh[i],
                    new_tex._num_faces_per_mesh[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(),
            ]
        )

        with self.assertRaises(ValueError):
            tex_mesh.extend(N=-1)

    def test_padded_to_packed(self):
        # 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.
        N = 2
        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)]

        num_faces_per_mesh = [f.shape[0] for f in faces_uvs_list]
        num_verts_per_mesh = [v.shape[0] for v in verts_uvs_list]
        tex = TexturesUV(
            maps=torch.ones((N, 16, 16, 3)),
            faces_uvs=faces_uvs_list,
            verts_uvs=verts_uvs_list,
        )

        # 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 = num_faces_per_mesh
        tex1._num_verts_per_mesh = num_verts_per_mesh
        verts_list = tex1.verts_uvs_list()
        verts_padded = tex1.verts_uvs_padded()

        faces_list = tex1.faces_uvs_list()
        faces_padded = tex1.faces_uvs_padded()

        for f1, f2 in zip(faces_list, faces_uvs_list):
            self.assertTrue((f1 == f2).all().item())

        for f1, f2 in zip(verts_list, verts_uvs_list):
            self.assertTrue((f1 == f2).all().item())

        self.assertTrue(faces_padded.shape == (2, 3, 3))
        self.assertTrue(verts_padded.shape == (2, 9, 2))

        # Case where num_faces_per_mesh is not set and faces_verts_uvs
        # are initialized with a padded tensor.
        tex2 = TexturesUV(
            maps=torch.ones((N, 16, 16, 3)),
            verts_uvs=verts_padded,
            faces_uvs=faces_padded,
        )
        faces_list = tex2.faces_uvs_list()
        verts_list = tex2.verts_uvs_list()

        for i, (f1, f2) in enumerate(zip(faces_list, faces_uvs_list)):
            n = num_faces_per_mesh[i]
            self.assertTrue((f1[:n] == f2).all().item())

        for i, (f1, f2) in enumerate(zip(verts_list, verts_uvs_list)):
            n = num_verts_per_mesh[i]
            self.assertTrue((f1[:n] == f2).all().item())

    def test_to(self):
        tex = TexturesUV(
            maps=torch.ones((5, 16, 16, 3)),
            faces_uvs=torch.randint(size=(5, 10, 3), high=15),
            verts_uvs=torch.rand(size=(5, 15, 2)),
        )
        device = torch.device("cuda:0")
        tex = tex.to(device)
        self.assertEqual(tex._faces_uvs_padded.device, device)
        self.assertEqual(tex._verts_uvs_padded.device, device)
        self.assertEqual(tex._maps_padded.device, device)

    def test_mesh_to(self):
        tex_cpu = TexturesUV(
            maps=torch.ones((5, 16, 16, 3)),
            faces_uvs=torch.randint(size=(5, 10, 3), high=15),
            verts_uvs=torch.rand(size=(5, 15, 2)),
        )
        verts = torch.rand(size=(5, 15, 3))
        faces = torch.randint(size=(5, 10, 3), high=15)
        mesh_cpu = Meshes(faces=faces, verts=verts, textures=tex_cpu)
        cpu = torch.device("cpu")
        device = torch.device("cuda:0")
        tex = mesh_cpu.to(device).textures
        self.assertEqual(tex._faces_uvs_padded.device, device)
        self.assertEqual(tex._verts_uvs_padded.device, device)
        self.assertEqual(tex._maps_padded.device, device)
        self.assertEqual(tex_cpu._verts_uvs_padded.device, cpu)

        self.assertEqual(tex_cpu.device, cpu)
        self.assertEqual(tex.device, device)

    def test_getitem(self):
        N = 5
        V = 20
        source = {
            "maps": torch.rand(size=(N, 1, 1, 3)),
            "faces_uvs": torch.randint(size=(N, 10, 3), high=V),
            "verts_uvs": torch.randn(size=(N, V, 2)),
        }
        tex = TexturesUV(
            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)

        tryindex(self, 2, tex, meshes, source)
        tryindex(self, slice(0, 2, 1), tex, meshes, source)
        index = torch.tensor([1, 0, 1, 0, 0], dtype=torch.bool)
        tryindex(self, index, tex, meshes, source)
        index = torch.tensor([0, 0, 0, 0, 0], dtype=torch.bool)
        tryindex(self, index, tex, meshes, source)
        index = torch.tensor([1, 2], dtype=torch.int64)
        tryindex(self, index, tex, meshes, source)
        tryindex(self, [2, 4], tex, meshes, source)

    def test_centers_for_image(self):
        maps = torch.rand(size=(1, 257, 129, 3))
        verts_uvs = torch.FloatTensor([[[0.25, 0.125], [0.5, 0.625], [0.5, 0.5]]])
        faces_uvs = torch.zeros(size=(1, 0, 3), dtype=torch.int64)
        tex = TexturesUV(maps=maps, faces_uvs=faces_uvs, verts_uvs=verts_uvs)

        expected = torch.FloatTensor([[32, 224], [64, 96], [64, 128]])
        self.assertClose(tex.centers_for_image(0), expected)


class TestRectanglePacking(TestCaseMixin, unittest.TestCase):
    def setUp(self) -> None:
        super().setUp()
        torch.manual_seed(42)

    def wrap_pack(self, sizes):
        """
        Call the pack_rectangles function, which we want to test,
        and return its outputs.
        Additionally makes some sanity checks on the output.
        """
        res = pack_rectangles(sizes)
        total = res.total_size
        self.assertGreaterEqual(total[0], 0)
        self.assertGreaterEqual(total[1], 0)
        mask = torch.zeros(total, dtype=torch.bool)
        seen_x_bound = False
        seen_y_bound = False
        for (in_x, in_y), loc in zip(sizes, res.locations):
            self.assertGreaterEqual(loc[0], 0)
            self.assertGreaterEqual(loc[1], 0)
            placed_x, placed_y = (in_y, in_x) if loc[2] else (in_x, in_y)
            upper_x = placed_x + loc[0]
            upper_y = placed_y + loc[1]
            self.assertGreaterEqual(total[0], upper_x)
            if total[0] == upper_x:
                seen_x_bound = True
            self.assertGreaterEqual(total[1], upper_y)
            if total[1] == upper_y:
                seen_y_bound = True
            already_taken = torch.sum(mask[loc[0] : upper_x, loc[1] : upper_y])
            self.assertEqual(already_taken, 0)
            mask[loc[0] : upper_x, loc[1] : upper_y] = 1
        self.assertTrue(seen_x_bound)
        self.assertTrue(seen_y_bound)

        self.assertTrue(torch.all(torch.sum(mask, dim=0, dtype=torch.int32) > 0))
        self.assertTrue(torch.all(torch.sum(mask, dim=1, dtype=torch.int32) > 0))
        return res

    def assert_bb(self, sizes, expected):
        """
        Apply the pack_rectangles function to sizes and verify the
        bounding box dimensions are expected.
        """
        self.assertSetEqual(set(self.wrap_pack(sizes).total_size), expected)

    def test_simple(self):
        self.assert_bb([(3, 4), (4, 3)], {6, 4})
        self.assert_bb([(2, 2), (2, 4), (2, 2)], {4, 4})

        # many squares
        self.assert_bb([(2, 2)] * 9, {2, 18})

        # One big square and many small ones.
        self.assert_bb([(3, 3)] + [(1, 1)] * 2, {3, 4})
        self.assert_bb([(3, 3)] + [(1, 1)] * 3, {3, 4})
        self.assert_bb([(3, 3)] + [(1, 1)] * 4, {3, 5})
        self.assert_bb([(3, 3)] + [(1, 1)] * 5, {3, 5})
        self.assert_bb([(1, 1)] * 6 + [(3, 3)], {3, 5})
        self.assert_bb([(3, 3)] + [(1, 1)] * 7, {3, 6})

        # many identical rectangles
        self.assert_bb([(7, 190)] * 4 + [(190, 7)] * 4, {190, 56})

        # require placing the flipped version of a rectangle
        self.assert_bb([(1, 100), (5, 96), (4, 5)], {100, 6})

    def test_random(self):
        for _ in range(5):
            vals = torch.randint(size=(20, 2), low=1, high=18)
            sizes = []
            for j in range(vals.shape[0]):
                sizes.append((int(vals[j, 0]), int(vals[j, 1])))
            self.wrap_pack(sizes)