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_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_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_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_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_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_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))
def test_initialize_lights_dimensions_fail(self): """ Color should have shape (N, 3) or (1, 3) """ with self.assertRaises(ValueError): DirectionalLights(ambient_color=torch.randn(10, 4)) with self.assertRaises(ValueError): DirectionalLights(direction=torch.randn(10, 4)) with self.assertRaises(ValueError): PointLights(ambient_color=torch.randn(10, 4)) with self.assertRaises(ValueError): PointLights(location=torch.randn(10, 4))
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_lights_accessor(self): d_light = DirectionalLights(ambient_color=((0.0, 0.0, 0.0), (1.0, 1.0, 1.0))) p_light = PointLights(ambient_color=((0.0, 0.0, 0.0), (1.0, 1.0, 1.0))) for light in [d_light, p_light]: # Update element color = (0.5, 0.5, 0.5) light[1].ambient_color = color self.assertClose(light.ambient_color[1], torch.tensor(color)) # Get item and get value l0 = light[0] self.assertClose(l0.ambient_color, torch.tensor((0.0, 0.0, 0.0)))
def test_initialize_lights_broadcast_fail(self): """ Batch dims have to be the same or 1. """ with self.assertRaises(ValueError): DirectionalLights(ambient_color=torch.randn(10, 3), diffuse_color=torch.randn(15, 3)) with self.assertRaises(ValueError): PointLights(ambient_color=torch.randn(10, 3), diffuse_color=torch.randn(15, 3))
def test_lights_clone_to(self): device = torch.device("cuda:0") cpu = torch.device("cpu") light = DirectionalLights() new_light = light.clone().to(device) keys = ["ambient_color", "diffuse_color", "specular_color", "direction"] for k in keys: prop = getattr(light, k) new_prop = getattr(new_light, k) self.assertTrue(prop.device == cpu) self.assertTrue(new_prop.device == device) self.assertSeparate(new_prop, prop) light = PointLights() new_light = light.clone().to(device) keys = ["ambient_color", "diffuse_color", "specular_color", "location"] for k in keys: prop = getattr(light, k) new_prop = getattr(new_light, k) self.assertTrue(prop.device == cpu) self.assertTrue(new_prop.device == device) self.assertSeparate(new_prop, prop)
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)
def test_initialize_lights_broadcast(self): light = DirectionalLights( ambient_color=torch.randn(10, 3), diffuse_color=torch.randn(1, 3), specular_color=torch.randn(1, 3), ) keys = ["ambient_color", "diffuse_color", "specular_color", "direction"] for k in keys: prop = getattr(light, k) self.assertTrue(prop.shape == (10, 3)) light = PointLights( ambient_color=torch.randn(10, 3), diffuse_color=torch.randn(1, 3), specular_color=torch.randn(1, 3), ) keys = ["ambient_color", "diffuse_color", "specular_color", "location"] for k in keys: prop = getattr(light, k) self.assertTrue(prop.shape == (10, 3))
def test_init_lights(self): """ Initialize Lights class with the default values. """ device = torch.device("cuda:0") light = DirectionalLights(device=device) keys = ["ambient_color", "diffuse_color", "specular_color", "direction"] for k in keys: prop = getattr(light, k) self.assertTrue(torch.is_tensor(prop)) self.assertTrue(prop.device == device) self.assertTrue(prop.shape == (1, 3)) light = PointLights(device=device) keys = ["ambient_color", "diffuse_color", "specular_color", "location"] for k in keys: prop = getattr(light, k) self.assertTrue(torch.is_tensor(prop)) self.assertTrue(prop.device == device) self.assertTrue(prop.shape == (1, 3))