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 forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
blend_params = kwargs.get("blend_params", self.blend_params)
texels = meshes.sample_textures(fragments)
images = hard_rgb_blend(texels, fragments, blend_params)
return images # (N, H, W, 3) RGBA image
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 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