def test_diffuse_directional_lights(self):
"""
Test with a single point where:
1) the normal and light direction are 45 degrees apart.
2) the normal and light direction are 90 degrees apart. The output
should be zero for this case
"""
color = torch.tensor([1, 1, 1], dtype=torch.float32)
direction = torch.tensor(
[0, 1 / np.sqrt(2), 1 / np.sqrt(2)], dtype=torch.float32
)
normals = torch.tensor([0, 0, 1], dtype=torch.float32)
normals = normals[None, None, :]
expected_output = torch.tensor(
[1 / np.sqrt(2), 1 / np.sqrt(2), 1 / np.sqrt(2)], dtype=torch.float32
)
expected_output = expected_output.view(1, 1, 3).repeat(3, 1, 1)
light = DirectionalLights(diffuse_color=color, direction=direction)
output_light = light.diffuse(normals=normals)
self.assertClose(output_light, expected_output)
# Change light direction to be 90 degrees apart from normal direction.
direction = torch.tensor([0, 1, 0], dtype=torch.float32)
light.direction = direction
expected_output = torch.zeros_like(expected_output)
output_light = light.diffuse(normals=normals)
self.assertClose(output_light, expected_output)
def test_diffuse_batched_packed(self):
"""
Test with a batch of 2 meshes each of which has faces on a single plane.
The normal and light direction are 45 degrees apart for the first mesh
and 90 degrees apart for the second mesh.
The points and normals are in the packed format i.e. no batch dimension.
"""
verts_packed = torch.rand((10, 3)) # points aren't used
faces_per_mesh = [6, 4]
mesh_to_vert_idx = [0] * faces_per_mesh[0] + [1] * faces_per_mesh[1]
mesh_to_vert_idx = torch.tensor(mesh_to_vert_idx, dtype=torch.int64)
color = torch.tensor([[1, 1, 1], [1, 1, 1]], dtype=torch.float32)
direction = torch.tensor(
[
[0, 1 / np.sqrt(2), 1 / np.sqrt(2)],
[0, 1, 0], # 90 degrees to normal so zero diffuse light
],
dtype=torch.float32,
)
normals = torch.tensor([[0, 0, 1], [0, 0, 1]], dtype=torch.float32)
expected_output = torch.zeros_like(verts_packed, dtype=torch.float32)
expected_output[:6, :] += 1 / np.sqrt(2)
expected_output[6:, :] = 0.0
lights = DirectionalLights(
diffuse_color=color[mesh_to_vert_idx, :],
direction=direction[mesh_to_vert_idx, :],
)
output_light = lights.diffuse(normals=normals[mesh_to_vert_idx, :])
self.assertClose(output_light, expected_output)
def test_diffuse_batched_arbitrary_input_dims(self):
"""
Test with a batch of inputs where shape of the input is mimicking the
shape in a shading function i.e. an interpolated normal per pixel for
top K faces per pixel.
"""
N, H, W, K = 16, 256, 256, 100
device = torch.device("cuda:0")
color = torch.tensor([1, 1, 1], dtype=torch.float32, device=device)
direction = torch.tensor(
[0, 1 / np.sqrt(2), 1 / np.sqrt(2)], dtype=torch.float32, device=device
)
normals = torch.tensor([0, 0, 1], dtype=torch.float32, device=device)
normals = normals.view(1, 1, 1, 1, 3).expand(N, H, W, K, -1)
direction = direction.view(1, 3)
color = color.view(1, 3)
expected_output = torch.tensor(
[1 / np.sqrt(2), 1 / np.sqrt(2), 1 / np.sqrt(2)],
dtype=torch.float32,
device=device,
)
expected_output = expected_output.view(1, 1, 1, 1, 3)
expected_output = expected_output.expand(N, H, W, K, -1)
lights = DirectionalLights(diffuse_color=color, direction=direction)
output_light = lights.diffuse(normals=normals)
self.assertClose(output_light, expected_output)
def test_diffuse_batched_broadcast_inputs(self):
"""
Test with a batch where each batch element has one point
where the normal and light direction are 45 degrees apart.
The color and direction are the same for each batch element.
"""
batch_size = 10
color = torch.tensor([1, 1, 1], dtype=torch.float32)
direction = torch.tensor(
[0, 1 / np.sqrt(2), 1 / np.sqrt(2)], dtype=torch.float32
)
normals = torch.tensor([0, 0, 1], dtype=torch.float32)
expected_out = torch.tensor(
[1 / np.sqrt(2), 1 / np.sqrt(2), 1 / np.sqrt(2)], dtype=torch.float32
)
# Reshape
normals = normals.view(-1, 1, 3).expand(batch_size, -1, -1)
expected_out = expected_out.view(-1, 1, 3).expand(batch_size, 1, 3)
# Don't expand the direction or color. Broadcasting should happen
# in the diffuse function.
direction = direction.view(1, 3)
color = color.view(1, 3)
lights = DirectionalLights(diffuse_color=color, direction=direction)
output_light = lights.diffuse(normals=normals)
self.assertClose(output_light, expected_out)
def test_diffuse_batched(self):
"""
Test with a batch where each batch element has one point
where the normal and light direction are 45 degrees apart.
"""
batch_size = 10
color = torch.tensor([1, 1, 1], dtype=torch.float32)
direction = torch.tensor(
[0, 1 / np.sqrt(2), 1 / np.sqrt(2)], dtype=torch.float32
)
normals = torch.tensor([0, 0, 1], dtype=torch.float32)
expected_out = torch.tensor(
[1 / np.sqrt(2), 1 / np.sqrt(2), 1 / np.sqrt(2)],
dtype=torch.float32,
)
# Reshape
direction = direction.view(-1, 3).expand(batch_size, -1)
normals = normals.view(-1, 1, 3).expand(batch_size, -1, -1)
color = color.view(-1, 3).expand(batch_size, -1)
expected_out = expected_out.view(-1, 1, 3).expand(batch_size, 1, 3)
lights = DirectionalLights(diffuse_color=color, direction=direction)
output_light = lights.diffuse(normals=normals)
self.assertTrue(torch.allclose(output_light, expected_out))