def test_specular_batched_packed(self): """ Test with a batch of 2 meshes each of which has faces on a single plane. The points and normals are in the packed format i.e. no batch dimension. """ 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, 0, 1]], dtype=torch.float32) direction = torch.tensor( [[-1 / np.sqrt(2), 1 / np.sqrt(2), 0], [-1, 1, 0]], dtype=torch.float32 ) camera_position = torch.tensor( [ [+1 / np.sqrt(2), 1 / np.sqrt(2), 0], [+1 / np.sqrt(2), -1 / np.sqrt(2), 0], ], dtype=torch.float32, ) points = torch.tensor([[0, 0, 0]], dtype=torch.float32) normals = torch.tensor([[0, 1, 0], [0, 1, 0]], dtype=torch.float32) expected_output = torch.zeros((10, 3), dtype=torch.float32) expected_output[:6, :] += 1.0 lights = DirectionalLights( specular_color=color[mesh_to_vert_idx, :], direction=direction[mesh_to_vert_idx, :], ) output_light = lights.specular( points=points.view(-1, 3).expand(10, -1), normals=normals.view(-1, 3)[mesh_to_vert_idx, :], camera_position=camera_position[mesh_to_vert_idx, :], shininess=10.0, ) self.assertClose(output_light, expected_output)
def test_specular_batched_broadcast_inputs(self): batch_size = 10 color = torch.tensor([1, 0, 1], dtype=torch.float32) direction = torch.tensor( [-1 / np.sqrt(2), 1 / np.sqrt(2), 0], dtype=torch.float32 ) camera_position = torch.tensor( [+1 / np.sqrt(2), 1 / np.sqrt(2), 0], dtype=torch.float32 ) points = torch.tensor([0, 0, 0], dtype=torch.float32) normals = torch.tensor([0, 1, 0], dtype=torch.float32) expected_out = torch.tensor([1.0, 0.0, 1.0], dtype=torch.float32) # Reshape normals = normals.view(1, 1, 3).expand(batch_size, -1, -1) points = points.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, color or camera_position. # These should be broadcasted in the specular function direction = direction.view(1, 3) camera_position = camera_position.view(1, 3) color = color.view(1, 3) lights = DirectionalLights(specular_color=color, direction=direction) output_light = lights.specular( points=points, normals=normals, camera_position=camera_position, shininess=torch.tensor(10), ) self.assertClose(output_light, expected_out)
def test_specular_batched(self): batch_size = 10 color = torch.tensor([1, 0, 1], dtype=torch.float32) direction = torch.tensor( [-1 / np.sqrt(2), 1 / np.sqrt(2), 0], dtype=torch.float32 ) camera_position = torch.tensor( [+1 / np.sqrt(2), 1 / np.sqrt(2), 0], dtype=torch.float32 ) points = torch.tensor([0, 0, 0], dtype=torch.float32) normals = torch.tensor([0, 1, 0], dtype=torch.float32) expected_out = torch.tensor([1.0, 0.0, 1.0], dtype=torch.float32) # Reshape direction = direction.view(1, 3).expand(batch_size, -1) camera_position = camera_position.view(1, 3).expand(batch_size, -1) normals = normals.view(1, 1, 3).expand(batch_size, -1, -1) points = points.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(specular_color=color, direction=direction) output_light = lights.specular( points=points, normals=normals, camera_position=camera_position, shininess=torch.tensor(10), ) self.assertTrue(torch.allclose(output_light, expected_out))
def test_specular_directional_lights(self): """ Specular highlights depend on the camera position as well as the light position/direction. Test with a single point where: 1) the normal and light direction are -45 degrees apart and the normal and camera position are +45 degrees apart. The reflected light ray will be perfectly aligned with the camera so the output is 1.0. 2) the normal and light direction are -45 degrees apart and the camera position is behind the point. The output should be zero for this case. """ color = torch.tensor([1, 0, 1], dtype=torch.float32) direction = torch.tensor( [-1 / np.sqrt(2), 1 / np.sqrt(2), 0], dtype=torch.float32 ) camera_position = torch.tensor( [+1 / np.sqrt(2), 1 / np.sqrt(2), 0], dtype=torch.float32 ) points = torch.tensor([0, 0, 0], dtype=torch.float32) normals = torch.tensor([0, 1, 0], dtype=torch.float32) expected_output = torch.tensor([1.0, 0.0, 1.0], dtype=torch.float32) expected_output = expected_output.view(1, 1, 3).repeat(3, 1, 1) lights = DirectionalLights(specular_color=color, direction=direction) output_light = lights.specular( points=points[None, None, :], normals=normals[None, None, :], camera_position=camera_position[None, :], shininess=torch.tensor(10), ) self.assertClose(output_light, expected_output) # Change camera position to be behind the point. camera_position = torch.tensor( [+1 / np.sqrt(2), -1 / np.sqrt(2), 0], dtype=torch.float32 ) expected_output = torch.zeros_like(expected_output) output_light = lights.specular( points=points[None, None, :], normals=normals[None, None, :], camera_position=camera_position[None, :], shininess=torch.tensor(10), ) self.assertClose(output_light, expected_output)
def test_specular_batched_arbitrary_input_dims(self): """ Test with a batch of inputs where shape of the input is mimicking the shape expected after rasterization i.e. a normal per pixel for top K faces per pixel. """ device = torch.device("cuda:0") N, H, W, K = 8, 128, 128, 100 color = torch.tensor([1, 0, 1], dtype=torch.float32, device=device) direction = torch.tensor( [-1 / np.sqrt(2), 1 / np.sqrt(2), 0], dtype=torch.float32 ) camera_position = torch.tensor( [+1 / np.sqrt(2), 1 / np.sqrt(2), 0], dtype=torch.float32 ) points = torch.tensor([0, 0, 0], dtype=torch.float32, device=device) normals = torch.tensor([0, 1, 0], dtype=torch.float32, device=device) points = points.view(1, 1, 1, 1, 3).expand(N, H, W, K, 3) normals = normals.view(1, 1, 1, 1, 3).expand(N, H, W, K, 3) direction = direction.view(1, 3) color = color.view(1, 3) camera_position = camera_position.view(1, 3) expected_output = torch.tensor( [1.0, 0.0, 1.0], 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(specular_color=color, direction=direction) output_light = lights.specular( points=points, normals=normals, camera_position=camera_position, shininess=10.0, ) self.assertClose(output_light, expected_output)