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