test_interpolate_face_attributes.py

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

import torch
from pytorch3d.ops.interp_face_attrs import (
    interpolate_face_attributes,
    interpolate_face_attributes_python,
)
from pytorch3d.renderer.mesh import TexturesVertex
from pytorch3d.renderer.mesh.rasterizer import Fragments
from pytorch3d.structures import Meshes

from .common_testing import get_random_cuda_device, TestCaseMixin


class TestInterpolateFaceAttributes(TestCaseMixin, unittest.TestCase):
    def _test_interp_face_attrs(self, interp_fun, device):
        pix_to_face = [0, 2, -1, 0, 1, -1]
        barycentric_coords = [
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
            [0.5, 0.5, 0.0],
            [0.8, 0.0, 0.2],
            [0.25, 0.5, 0.25],
        ]
        face_attrs = [
            [[1, 2], [3, 4], [5, 6]],
            [[7, 8], [9, 10], [11, 12]],
            [[13, 14], [15, 16], [17, 18]],
        ]
        pix_attrs = [
            [1, 2],
            [15, 16],
            [0, 0],
            [2, 3],
            [0.8 * 7 + 0.2 * 11, 0.8 * 8 + 0.2 * 12],
            [0, 0],
        ]
        N, H, W, K, D = 1, 2, 1, 3, 2
        pix_to_face = torch.tensor(pix_to_face, dtype=torch.int64, device=device)
        pix_to_face = pix_to_face.view(N, H, W, K)
        barycentric_coords = torch.tensor(
            barycentric_coords, dtype=torch.float32, device=device
        )
        barycentric_coords = barycentric_coords.view(N, H, W, K, 3)
        face_attrs = torch.tensor(face_attrs, dtype=torch.float32, device=device)
        pix_attrs = torch.tensor(pix_attrs, dtype=torch.float32, device=device)
        pix_attrs = pix_attrs.view(N, H, W, K, D)

        args = (pix_to_face, barycentric_coords, face_attrs)
        pix_attrs_actual = interp_fun(*args)
        self.assertClose(pix_attrs_actual, pix_attrs)

    def test_python(self):
        device = torch.device("cuda:0")
        self._test_interp_face_attrs(interpolate_face_attributes_python, device)

    def test_cuda(self):
        device = torch.device("cuda:0")
        self._test_interp_face_attrs(interpolate_face_attributes, device)

    def test_python_vs_cuda(self):
        N, H, W, K = 2, 32, 32, 5
        F = 1000
        D = 3
        device = get_random_cuda_device()
        torch.manual_seed(598)
        pix_to_face = torch.randint(-F, F, (N, H, W, K), device=device)
        barycentric_coords = torch.randn(
            N, H, W, K, 3, device=device, requires_grad=True
        )
        face_attrs = torch.randn(F, 3, D, device=device, requires_grad=True)
        grad_pix_attrs = torch.randn(N, H, W, K, D, device=device)
        args = (pix_to_face, barycentric_coords, face_attrs)

        # Run the python version
        pix_attrs_py = interpolate_face_attributes_python(*args)
        pix_attrs_py.backward(gradient=grad_pix_attrs)
        grad_bary_py = barycentric_coords.grad.clone()
        grad_face_attrs_py = face_attrs.grad.clone()

        # Clear gradients
        barycentric_coords.grad.zero_()
        face_attrs.grad.zero_()

        # Run the CUDA version
        pix_attrs_cu = interpolate_face_attributes(*args)
        pix_attrs_cu.backward(gradient=grad_pix_attrs)
        grad_bary_cu = barycentric_coords.grad.clone()
        grad_face_attrs_cu = face_attrs.grad.clone()

        # Check they are the same
        self.assertClose(pix_attrs_py, pix_attrs_cu, rtol=2e-3)
        self.assertClose(grad_bary_py, grad_bary_cu, rtol=1e-4)
        self.assertClose(grad_face_attrs_py, grad_face_attrs_cu, rtol=1e-3)

    def test_interpolate_attributes(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
        )
        tex = TexturesVertex(verts_features=vert_tex[None, :])
        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),
        )

        verts_features_packed = mesh.textures.verts_features_packed()
        faces_verts_features = verts_features_packed[mesh.faces_packed()]

        texels = interpolate_face_attributes(
            fragments.pix_to_face, fragments.bary_coords, faces_verts_features
        )
        self.assertTrue(torch.allclose(texels, expected_vals[None, :]))

    def test_interpolate_attributes_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,
        )
        tex = TexturesVertex(verts_features=vert_tex[None, :])
        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,
        )
        verts_features_packed = mesh.textures.verts_features_packed()
        faces_verts_features = verts_features_packed[mesh.faces_packed()]

        texels = interpolate_face_attributes(
            fragments.pix_to_face, fragments.bary_coords, faces_verts_features
        )
        texels.sum().backward()
        self.assertTrue(hasattr(vert_tex, "grad"))
        self.assertTrue(torch.allclose(vert_tex.grad, grad_vert_tex[None, :]))

    def test_interpolate_face_attributes_fail(self):
        # 1. A face can only have 3 verts
        #   i.e. face_attributes must have shape (F, 3, D)
        face_attributes = torch.ones(1, 4, 3)
        pix_to_face = torch.ones((1, 1, 1, 1))
        fragments = Fragments(
            pix_to_face=pix_to_face,
            bary_coords=pix_to_face[..., None].expand(-1, -1, -1, -1, 3),
            zbuf=pix_to_face,
            dists=pix_to_face,
        )
        with self.assertRaises(ValueError):
            interpolate_face_attributes(
                fragments.pix_to_face, fragments.bary_coords, face_attributes
            )

        # 2. pix_to_face must have shape (N, H, W, K)
        pix_to_face = torch.ones((1, 1, 1, 1, 3))
        fragments = Fragments(
            pix_to_face=pix_to_face,
            bary_coords=pix_to_face,
            zbuf=pix_to_face,
            dists=pix_to_face,
        )
        with self.assertRaises(ValueError):
            interpolate_face_attributes(
                fragments.pix_to_face, fragments.bary_coords, face_attributes
            )