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