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
)
def test_softmax_rgb_blend(self):
# Create dummy outputs of rasterization simulating a cube in the center
# of the image with surrounding padded values.
N, S, K = 1, 8, 2
device = torch.device("cuda")
pix_to_face = torch.full((N, S, S, K),
fill_value=-1,
dtype=torch.int64,
device=device)
h = int(S / 2)
pix_to_face_full = torch.randint(size=(N, h, h, K),
low=0,
high=100,
device=device)
s = int(S / 4)
e = int(0.75 * S)
pix_to_face[:, s:e, s:e, :] = pix_to_face_full
empty = torch.tensor([], device=device)
random_sign_flip = torch.rand((N, S, S, K), device=device)
random_sign_flip[random_sign_flip > 0.5] *= -1.0
zbuf1 = torch.randn(size=(N, S, S, K), device=device)
# randomly flip the sign of the distance
# (-) means inside triangle, (+) means outside triangle.
dists1 = torch.randn(size=(N, S, S, K),
device=device) * random_sign_flip
dists2 = dists1.clone()
zbuf2 = zbuf1.clone()
dists1.requires_grad = True
dists2.requires_grad = True
colors = torch.randn((N, S, S, K, 3), device=device)
fragments1 = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty, # dummy
zbuf=zbuf1,
dists=dists1,
)
fragments2 = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty, # dummy
zbuf=zbuf2,
dists=dists2,
)
blend_params = BlendParams(sigma=1e-3)
args1 = (colors, fragments1, blend_params)
args2 = (colors, fragments2, blend_params)
self._compare_impls(
softmax_rgb_blend,
softmax_blend_naive,
args1,
args2,
dists1,
dists2,
compare_grads=True,
)
def test_softmax_rgb_blend(self):
# Create dummy outputs of rasterization simulating a cube in the centre
# of the image with surrounding padded values.
N, S, K = 1, 8, 2
pix_to_face = -torch.ones((N, S, S, K), dtype=torch.int64)
h = int(S / 2)
pix_to_face_full = torch.randint(size=(N, h, h, K), low=0, high=100)
s = int(S / 4)
e = int(0.75 * S)
pix_to_face[:, s:e, s:e, :] = pix_to_face_full
bary_coords = torch.ones((N, S, S, K, 3))
random_sign_flip = torch.rand((N, S, S, K))
random_sign_flip[random_sign_flip > 0.5] *= -1.0
zbuf1 = torch.randn(size=(N, S, S, K))
# randomly flip the sign of the distance
# (-) means inside triangle, (+) means outside triangle.
dists1 = torch.randn(size=(N, S, S, K)) * random_sign_flip
dists2 = dists1.clone()
zbuf2 = zbuf1.clone()
dists1.requires_grad = True
dists2.requires_grad = True
zbuf1.requires_grad = True
zbuf2.requires_grad = True
colors = torch.randn_like(bary_coords)
fragments1 = Fragments(
pix_to_face=pix_to_face,
bary_coords=bary_coords, # dummy
zbuf=zbuf1,
dists=dists1,
)
fragments2 = Fragments(
pix_to_face=pix_to_face,
bary_coords=bary_coords, # dummy
zbuf=zbuf2,
dists=dists2,
)
blend_params = BlendParams(sigma=1e-1)
images = softmax_rgb_blend(colors, fragments1, blend_params)
images_naive = softmax_blend_naive(colors, fragments2, blend_params)
self.assertTrue(torch.allclose(images, images_naive))
# Check gradients.
images.sum().backward()
self.assertTrue(hasattr(dists1, "grad"))
self.assertTrue(hasattr(zbuf1, "grad"))
images_naive.sum().backward()
self.assertTrue(hasattr(dists2, "grad"))
self.assertTrue(hasattr(zbuf2, "grad"))
self.assertTrue(torch.allclose(dists1.grad, dists2.grad, atol=2e-5))
self.assertTrue(torch.allclose(zbuf1.grad, zbuf2.grad, atol=2e-5))
# Helpful comments below.# Helpful comments below.
def test_sigmoid_alpha_blend_python(self):
"""
Test outputs of python tensorised function and python loop
"""
# Create dummy outputs of rasterization
torch.manual_seed(231)
F = 32 # number of faces in the mesh
# The python loop version is really slow so only using small input sizes.
N, S, K = 2, 10, 5
device = torch.device('cuda')
pix_to_face = torch.randint(F + 1, size=(N, S, S, K),
device=device) - 1
colors = torch.randn((N, S, S, K, 3), device=device)
empty = torch.tensor([], device=device)
# # randomly flip the sign of the distance
# # (-) means inside triangle, (+) means outside triangle.
random_sign_flip = torch.rand((N, S, S, K))
random_sign_flip[random_sign_flip > 0.5] *= -1.0
dists1 = torch.randn(size=(N, S, S, K),
requires_grad=True,
device=device)
dists2 = dists1.detach().clone()
dists2.requires_grad = True
fragments1 = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty, # dummy
zbuf=empty, # dummy
dists=dists1,
)
fragments2 = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty, # dummy
zbuf=empty, # dummy
dists=dists2,
)
blend_params = BlendParams(sigma=1e-2)
args1 = (colors, fragments1, blend_params)
args2 = (colors, fragments2, blend_params)
self._compare_impls(
sigmoid_alpha_blend,
sigmoid_blend_naive_loop,
args1,
args2,
dists1,
dists2,
compare_grads=True,
)
def test_sigmoid_alpha_blend(self):
"""
Test outputs of sigmoid alpha blend tensorised function match those of
the naive iterative version. Also check gradients match.
"""
# Create dummy outputs of rasterization simulating a cube in the centre
# of the image with surrounding padded values.
N, S, K = 1, 8, 2
pix_to_face = -torch.ones((N, S, S, K), dtype=torch.int64)
h = int(S / 2)
pix_to_face_full = torch.randint(size=(N, h, h, K), low=0, high=100)
s = int(S / 4)
e = int(0.75 * S)
pix_to_face[:, s:e, s:e, :] = pix_to_face_full
bary_coords = torch.ones((N, S, S, K, 3))
# randomly flip the sign of the distance
# (-) means inside triangle, (+) means outside triangle.
random_sign_flip = torch.rand((N, S, S, K))
random_sign_flip[random_sign_flip > 0.5] *= -1.0
dists = torch.randn(size=(N, S, S, K))
dists1 = dists * random_sign_flip
dists2 = dists1.clone()
dists1.requires_grad = True
dists2.requires_grad = True
colors = torch.randn_like(bary_coords)
fragments1 = Fragments(
pix_to_face=pix_to_face,
bary_coords=bary_coords, # dummy
zbuf=pix_to_face, # dummy
dists=dists1,
)
fragments2 = Fragments(
pix_to_face=pix_to_face,
bary_coords=bary_coords, # dummy
zbuf=pix_to_face, # dummy
dists=dists2,
)
blend_params = BlendParams(sigma=2e-1)
images = sigmoid_alpha_blend(colors, fragments1, blend_params)
images_naive = sigmoid_blend_naive(colors, fragments2, blend_params)
self.assertTrue(torch.allclose(images, images_naive))
torch.manual_seed(231)
images.sum().backward()
self.assertTrue(hasattr(dists1, "grad"))
images_naive.sum().backward()
self.assertTrue(hasattr(dists2, "grad"))
self.assertTrue(torch.allclose(dists1.grad, dists2.grad, rtol=1e-5))
def test_sigmoid_alpha_blend_python(self):
"""
Test outputs of python tensorised function and python loop
"""
# Create dummy outputs of rasterization
torch.manual_seed(231)
F = 32 # number of faces in the mesh
# The python loop version is really slow so only using small input sizes.
N, S, K = 1, 4, 1
device = torch.device("cuda")
pix_to_face = torch.randint(low=-1,
high=F,
size=(N, S, S, K),
device=device)
colors = torch.randn((N, S, S, K, 3), device=device)
empty = torch.tensor([], device=device)
dists1 = torch.randn(size=(N, S, S, K), device=device)
dists2 = dists1.clone()
dists1.requires_grad = True
dists2.requires_grad = True
fragments1 = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty, # dummy
zbuf=empty, # dummy
dists=dists1,
)
fragments2 = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty, # dummy
zbuf=empty, # dummy
dists=dists2,
)
blend_params = BlendParams(sigma=1e-2)
args1 = (colors, fragments1, blend_params)
args2 = (colors, fragments2, blend_params)
self._compare_impls(
sigmoid_alpha_blend,
sigmoid_alpha_blend_vectorized,
args1,
args2,
dists1,
dists2,
compare_grads=True,
)
def test_sigmoid_alpha_blend_manual_gradients(self):
# Create dummy outputs of rasterization
torch.manual_seed(231)
F = 32 # number of faces in the mesh
# The python loop version is really slow so only using small input sizes.
N, S, K = 2, 3, 2
device = torch.device("cuda")
pix_to_face = torch.randint(F + 1, size=(N, S, S, K),
device=device) - 1
colors = torch.randn((N, S, S, K, 3), device=device)
empty = torch.tensor([], device=device)
# # randomly flip the sign of the distance
# # (-) means inside triangle, (+) means outside triangle.
random_sign_flip = torch.rand((N, S, S, K))
random_sign_flip[random_sign_flip > 0.5] *= -1.0
dists = torch.randn(size=(N, S, S, K),
requires_grad=True,
device=device)
fragments = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty, # dummy
zbuf=empty, # dummy
dists=dists,
)
blend_params = BlendParams(sigma=1e-3)
pix_cols = sigmoid_blend_naive_loop(colors, fragments, blend_params)
grad_out = torch.randn_like(pix_cols)
# Backward pass
pix_cols.backward(grad_out)
grad_dists = sigmoid_blend_naive_loop_backward(grad_out, pix_cols,
fragments, blend_params)
self.assertTrue(torch.allclose(dists.grad, grad_dists, atol=1e-7))
def test_hard_rgb_blend(self):
N, H, W, K = 5, 10, 10, 20
pix_to_face = torch.randint(low=-1, high=100, size=(N, H, W, K))
bary_coords = torch.ones((N, H, W, K, 3))
fragments = Fragments(
pix_to_face=pix_to_face,
bary_coords=bary_coords,
zbuf=pix_to_face, # dummy
dists=pix_to_face, # dummy
)
colors = torch.randn((N, H, W, K, 3))
blend_params = BlendParams(1e-4, 1e-4, (0.5, 0.5, 1))
images = hard_rgb_blend(colors, fragments, blend_params)
# Examine if the foreground colors are correct.
is_foreground = pix_to_face[..., 0] >= 0
self.assertClose(images[is_foreground][:, :3],
colors[is_foreground][..., 0, :])
# Examine if the background colors are correct.
for i in range(3): # i.e. RGB
channel_color = blend_params.background_color[i]
self.assertTrue(images[~is_foreground][...,
i].eq(channel_color).all())
# Examine the alpha channel
self.assertClose(images[..., 3], (pix_to_face[..., 0] >= 0).float())
def test_cameras_check(self):
verts = torch.tensor([[-1, -1, 0], [1, -1, 1], [1, 1, 0], [-1, 1, 1]],
dtype=torch.float32)
faces = torch.tensor([[0, 1, 2], [2, 3, 0]], dtype=torch.int64)
meshes = Meshes(verts=[verts], faces=[faces])
pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
barycentric_coords = torch.tensor([[0.1, 0.2, 0.7], [0.3, 0.5, 0.2]],
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),
)
shader_classes = [
HardFlatShader,
HardGouraudShader,
HardPhongShader,
SoftPhongShader,
]
for shader_class in shader_classes:
shader = shader_class()
with self.assertRaises(ValueError):
shader(fragments, meshes)
def forward(self, meshes_world, **kwargs) -> torch.Tensor:
"""
Render a batch of images from a batch of meshes by rasterizing and then shading.
NOTE: If the blur radius for rasterizaiton is > 0.0, some pixels can have one or more barycentric coordinates lying outside the range [0, 1].
For a pixel with out of bounds barycentric coordinates with respect to a face f, clipping is required before interpolating the texture uv coordinates and z buffer so that the colors and depths are limited to the range for the corresponding face.
"""
fragments = self.rasterizer(meshes_world, **kwargs)
raster_setting = kwargs.get("raster_settings",
self.rasterizer.raster_settings)
if raster_setting.blur_radius > 0.0:
# TODO: potentially move barycenteric clipping to the rasterizer.
# If no downstream functions requiers upclipped values.
# This will avoid uncessary re-interpolation of the z buffer.
clipped_bary_coords = _clip_barycentric_coordinates(
fragments.bary_coords)
clipped_zbuf = _interpolate_zbuf(fragments.pix_to_face,
clipped_bary_coords, meshes_world)
fragments = Fragments(
bary_coords=clipped_bary_coords,
zbuf=clipped_zbuf,
dists=fragments.dists,
pix_to_face=fragments.pix_to_face,
)
images = self.shader(fragments, meshes_world, **kwargs)
depth = fragments.zbuf
return images, depth
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_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_interpolate_attributes(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)
tex = Textures(verts_rgb=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),
)
texels = interpolate_vertex_colors(fragments, mesh)
self.assertTrue(torch.allclose(texels, expected_vals[None, :]))
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_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 raster_fn():
fragments = rasterizer(sphere_meshes)
# Clip bary and reinterpolate
clipped_bary_coords = _clip_barycentric_coordinates(
fragments.bary_coords)
clipped_zbuf = _interpolate_zbuf(fragments.pix_to_face,
clipped_bary_coords, sphere_meshes)
fragments = Fragments(
bary_coords=clipped_bary_coords,
zbuf=clipped_zbuf,
dists=fragments.dists,
pix_to_face=fragments.pix_to_face,
)
torch.cuda.synchronize()
def bm_softmax_blending(
num_meshes: int = 16,
image_size: int = 128,
faces_per_pixel: int = 100,
device: str = 'cpu',
):
if torch.cuda.is_available() and 'cuda:' in device:
# If a device other than the default is used, set the device explicity.
torch.cuda.set_device(device)
device = torch.device(device)
torch.manual_seed(231)
# Create dummy outputs of rasterization
N, S, K = num_meshes, image_size, faces_per_pixel
F = 32 # num faces in the mesh
pix_to_face = torch.randint(F + 1, size=(N, S, S, K),
device=device) - 1
colors = torch.randn((N, S, S, K, 3), device=device)
empty = torch.tensor([], device=device)
# # randomly flip the sign of the distance
# # (-) means inside triangle, (+) means outside triangle.
random_sign_flip = torch.rand((N, S, S, K), device=device)
random_sign_flip[random_sign_flip > 0.5] *= -1.0
dists1 = torch.randn(size=(N, S, S, K),
requires_grad=True,
device=device)
zbuf = torch.randn(size=(N, S, S, K),
requires_grad=True,
device=device)
fragments = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty,
zbuf=zbuf,
dists=dists1, # dummy
)
blend_params = BlendParams(sigma=1e-3)
torch.cuda.synchronize()
def fn():
# test forward and backward pass
images = softmax_rgb_blend(colors, fragments, blend_params)
images.sum().backward()
torch.cuda.synchronize()
return fn
def bm_softmax_blending(
num_meshes: int = 16,
image_size: int = 128,
faces_per_pixel: int = 100,
device: str = "cpu",
backend: str = "pytorch",
):
if torch.cuda.is_available() and "cuda:" in device:
# If a device other than the default is used, set the device explicity.
torch.cuda.set_device(device)
device = torch.device(device)
torch.manual_seed(231)
# Create dummy outputs of rasterization
N, S, K = num_meshes, image_size, faces_per_pixel
F = 32 # num faces in the mesh
pix_to_face = torch.randint(low=-1,
high=F + 1,
size=(N, S, S, K),
device=device)
colors = torch.randn((N, S, S, K, 3), device=device)
empty = torch.tensor([], device=device)
dists1 = torch.randn(size=(N, S, S, K),
requires_grad=True,
device=device)
zbuf = torch.randn(size=(N, S, S, K),
requires_grad=True,
device=device)
fragments = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty,
zbuf=zbuf,
dists=dists1 # dummy
)
blend_params = BlendParams(sigma=1e-3)
torch.cuda.synchronize()
def fn():
# test forward and backward pass
images = softmax_rgb_blend(colors, fragments, blend_params)
images.sum().backward()
torch.cuda.synchronize()
return fn
def rasterize(R, T, meshes, rasterizer, blur_radius=0):
# It will automatically update the camera settings -> R, T in rasterizer.camera
fragments = rasterizer(meshes, R=R, T=T)
# Copy from pytorch3D source code, try if it is necessary to do gradient decent
if blur_radius > 0.0:
clipped_bary_coords = utils._clip_barycentric_coordinates(
fragments.bary_coords)
clipped_zbuf = utils._interpolate_zbuf(fragments.pix_to_face,
clipped_bary_coords, meshes)
fragments = Fragments(
bary_coords=clipped_bary_coords,
zbuf=clipped_zbuf,
dists=fragments.dists,
pix_to_face=fragments.pix_to_face,
)
return fragments
def _rasterize(self, meshes, image_size, bin_size, blur):
"""
Simple wrapper around the rasterize function to return
the fragment data.
"""
face_idxs, zbuf, bary_coords, pix_dists = rasterize_meshes(
meshes,
image_size,
blur,
faces_per_pixel=1,
bin_size=bin_size,
)
return Fragments(
pix_to_face=face_idxs,
zbuf=zbuf,
bary_coords=bary_coords,
dists=pix_dists,
)
def test_hard_rgb_blend(self):
N, H, W, K = 5, 10, 10, 20
pix_to_face = torch.ones((N, H, W, K))
bary_coords = torch.ones((N, H, W, K, 3))
fragments = Fragments(
pix_to_face=pix_to_face,
bary_coords=bary_coords,
zbuf=pix_to_face, # dummy
dists=pix_to_face, # dummy
)
colors = bary_coords.clone()
top_k = torch.randn((K, 3))
colors[..., :, :] = top_k
images = hard_rgb_blend(colors, fragments)
expected_vals = torch.ones((N, H, W, 4))
pix_cols = torch.ones_like(expected_vals[..., :3]) * top_k[0, :]
expected_vals[..., :3] = pix_cols
self.assertTrue(torch.allclose(images, expected_vals))
def bm_sigmoid_alpha_blending(
num_meshes: int = 16,
image_size: int = 128,
faces_per_pixel: int = 100,
device="cuda",
backend: str = "pytorch",
):
device = torch.device(device)
torch.manual_seed(231)
# Create dummy outputs of rasterization
N, S, K = num_meshes, image_size, faces_per_pixel
F = 32 # num faces in the mesh
pix_to_face = torch.randint(low=-1,
high=F + 1,
size=(N, S, S, K),
device=device)
colors = torch.randn((N, S, S, K, 3), device=device)
empty = torch.tensor([], device=device)
dists1 = torch.randn(size=(N, S, S, K),
requires_grad=True,
device=device)
fragments = Fragments(
pix_to_face=pix_to_face,
bary_coords=empty, # dummy
zbuf=empty, # dummy
dists=dists1,
)
blend_params = BlendParams(sigma=1e-3)
blend_fn = (sigmoid_alpha_blend_vectorized
if backend == "pytorch" else sigmoid_alpha_blend)
torch.cuda.synchronize()
def fn():
# test forward and backward pass
images = blend_fn(colors, fragments, blend_params)
images.sum().backward()
torch.cuda.synchronize()
return fn
def test_cpu(self):
"""
Test that the output of rendering non square images
gives the same result as square images. i.e. the
dists, zbuf, bary are all the same for the square
region which is present in both images.
In this test we compare between the naive C++ implementation
and the naive python implementation as the Coarse/Fine
method is not fully implemented in C++
"""
# Test both when (W > H) and (H > W).
# Using smaller image sizes here as the Python rasterizer is really slow.
image_sizes = [(32, 64), (64, 32), (60, 110)]
devices = ["cpu"]
blurs = [0.0, 0.001]
batch_sizes = [1]
test_cases = product(image_sizes, blurs, devices, batch_sizes)
for image_size, blur, device, batch_size in test_cases:
# Initialize the verts grad tensor and the meshes objects
verts_nonsq_naive, meshes_nonsq_naive = self._clone_mesh(
verts0, faces0, device, batch_size)
verts_nonsq_python, meshes_nonsq_python = self._clone_mesh(
verts0, faces0, device, batch_size)
# Compare Naive CPU with Python as Coarse/Fine rasteriztation
# is not implemented for CPU
fragments_naive = self._rasterize(meshes_nonsq_naive,
image_size,
bin_size=0,
blur=blur)
face_idxs, zbuf, bary_coords, pix_dists = rasterize_meshes_python(
meshes_nonsq_python,
image_size,
blur,
faces_per_pixel=1,
)
fragments_python = Fragments(
pix_to_face=face_idxs,
zbuf=zbuf,
bary_coords=bary_coords,
dists=pix_dists,
)
# Save debug images if DEBUG is set to true at the top of the file.
_save_debug_image(fragments_naive.pix_to_face, image_size, 0, blur)
_save_debug_image(fragments_python.pix_to_face, image_size,
"python", blur)
# List of non square outputs to compare with the square output
nonsq_fragment_gradtensor_list = [
(fragments_naive, verts_nonsq_naive, "naive"),
(fragments_python, verts_nonsq_python, "python"),
]
self._compare_square_with_nonsq(
image_size,
blur,
device,
verts0,
faces0,
nonsq_fragment_gradtensor_list,
batch_size,
)