def test_blend_params(self):
"""Test color parameter of BlendParams().
Assert passed value overrides default value.
"""
bp_default = BlendParams()
bp_new = BlendParams(background_color=(0.5, 0.5, 0.5))
self.assertEqual(bp_new.background_color, (0.5, 0.5, 0.5))
self.assertEqual(bp_default.background_color, (1.0, 1.0, 1.0))
def soft_feature_blending(colors,
fragments,
blend_params=None,
znear: float = 0.01,
zfar: float = 10) -> torch.Tensor:
"""
Returns:
Rendered features: (N, H, W, F)
"""
blend_params = BlendParams() if blend_params is None else blend_params
eps = 1e-10 # Weight for background
N, H, W, K = fragments.pix_to_face.shape
# number of feature channels
# C = colors.shape[-1]
mask = fragments.pix_to_face >= 0
prob_map = torch.sigmoid(-fragments.dists / blend_params.sigma) * mask
z_inv = (zfar - fragments.zbuf) / (zfar - znear) * mask
# pyre-fixme[16]: `Tuple` has no attribute `values`.
# pyre-fixme[6]: Expected `Tensor` for 1st param but got `float`.
z_inv_max = torch.max(z_inv, dim=-1).values[..., None].clamp(min=eps)
# pyre-fixme[6]: Expected `Tensor` for 1st param but got `float`.
z_prob = torch.exp((z_inv - z_inv_max) / blend_params.gamma)
z_prob = z_prob / (z_prob.sum(-1, keepdim=True).clamp(eps, 1))
# For each face, compute the soft assignment score by taking into account xy dist
# and relative z distance
# (among different faces assigned to specific xy pixel location)
colored = 1 - torch.prod(1 - (prob_map * z_prob).unsqueeze(-1) * colors,
dim=-2)
alpha = 1 - torch.prod((1.0 - prob_map), dim=-1)
return torch.cat([colored, alpha.unsqueeze(-1)], -1)
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_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 setup(self, device):
R, T = look_at_view_transform(self.viewpoint_distance,
self.viewpoint_elevation,
self.viewpoint_azimuth,
device=device)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
raster_settings = RasterizationSettings(
image_size=self.opt.fast_image_size,
blur_radius=self.opt.raster_blur_radius,
faces_per_pixel=self.opt.raster_faces_per_pixel,
)
rasterizer = MeshRasterizer(cameras=cameras,
raster_settings=raster_settings)
lights = PointLights(device=device,
location=[self.opt.lights_location])
lights = DirectionalLights(device=device,
direction=[self.opt.lights_direction])
shader = SoftPhongShader(
device=device,
cameras=cameras,
lights=lights,
blend_params=BlendParams(
self.opt.blend_params_sigma,
self.opt.blend_params_gamma,
self.opt.blend_params_background_color,
),
)
self.renderer = MeshRenderer(
rasterizer=rasterizer,
shader=shader,
)
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 render_torch(self,
verts,
faces,
rgb,
bcg_color=(1., 1., 1.),
get_depth=False,
get_alpha=False):
# b, h, w = grid_3d.shape[:3]
b = verts.size(0)
textures = TexturesVertex(verts_features=rgb.view(b, -1, 3))
mesh = Meshes(verts=verts, faces=faces, textures=textures)
fragments = self.rasterizer_torch(mesh)
texels = mesh.sample_textures(fragments)
materials = Materials(device=verts.device)
blend_params = BlendParams(background_color=bcg_color)
images = hard_rgb_blend(texels, fragments, blend_params)
images = images[..., :3].permute(0, 3, 1, 2)
out = (images, )
if get_depth:
depth = fragments.zbuf[..., 0]
mask = (depth == -1.0).float()
max_depth = self.max_depth + 0.5 * (self.max_depth -
self.min_depth)
depth = mask * max_depth * torch.ones_like(depth) + (1 -
mask) * depth
out = out + (depth, )
if get_alpha:
colors = torch.ones_like(fragments.bary_coords)
blend_params = BlendParams(sigma=1e-2,
gamma=1e-4,
background_color=(1., 1., 1.))
alpha = sigmoid_alpha_blend(colors, fragments, blend_params)[...,
-1]
out = tuple(out) + (alpha, )
if len(out) == 1:
out = out[0]
return out
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 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 __init__(self,
background_color: Optional[Union[Tuple, List,
torch.Tensor]] = None,
device="cpu",
cameras=None,
lights=None,
materials=None,
blend_params=None):
super().__init__()
self.background_color = background_color
self.lights = lights if lights is not None else PointLights(
device=device)
self.materials = (materials if materials is not None else Materials(
device=device))
self.cameras = cameras
self.blend_params = blend_params if blend_params is not None else BlendParams(
)
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
