def test_jit_trace(self, device): @torch.jit.script def op_script(input, map1, map2): return kornia.remap(input, map1, map2) # 1. Trace op batch_size, channels, height, width = 1, 1, 3, 4 img = torch.ones(batch_size, channels, height, width).to(device) grid = kornia.utils.create_meshgrid( height, width, normalized_coordinates=False).to(device) grid += 1. # apply some shift input_tuple = (img, grid[..., 0], grid[..., 1]) op_traced = torch.jit.trace(op_script, input_tuple) # 2. Generate different input batch_size, channels, height, width = 2, 2, 2, 5 img = torch.ones(batch_size, channels, height, width).to(device) grid = kornia.utils.create_meshgrid( height, width, normalized_coordinates=False).to(device) grid += 2. # apply some shift # 3. Apply to different input input_tuple = (img, grid[..., 0], grid[..., 1]) actual = op_script(*input_tuple) expected = kornia.remap(*input_tuple) assert_allclose(actual, expected)
def test_jit(self, device, dtype): @torch.jit.script def op_script(input, map1, map2): return kornia.remap(input, map1, map2) batch_size, channels, height, width = 1, 1, 3, 4 img = torch.ones(batch_size, channels, height, width, device=device, dtype=dtype) grid = kornia.utils.create_meshgrid(height, width, normalized_coordinates=False, device=device).to(dtype) grid += 1. # apply some shift input = ( img, grid[..., 0], grid[..., 1], ) actual = op_script(*input) expected = kornia.remap(*input) assert_allclose(actual, expected, rtol=1e-4, atol=1e-4)
def test_shift_batch(self, device): height, width = 3, 4 inp = torch.tensor([[[ [1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.], ]]]).repeat(2, 1, 1, 1).to(device) expected = torch.tensor([[[ [1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.], ]], [[ [1., 1., 1., 1.], [1., 1., 1., 1.], [0., 0., 0., 0.], ]]]).to(device) # generate a batch of grids grid = kornia.utils.create_meshgrid( height, width, normalized_coordinates=False).to(device) grid = grid.repeat(2, 1, 1, 1) grid[0, ..., 0] += 1. # apply shift in the x direction grid[1, ..., 1] += 1. # apply shift in the y direction input_warped = kornia.remap(inp, grid[..., 0], grid[..., 1], align_corners=True) assert_allclose(input_warped, expected)
def test_shift_batch_broadcast(self, device, dtype): height, width = 3, 4 inp = torch.tensor([[[ [1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.], ]]], device=device, dtype=dtype).repeat(2, 1, 1, 1) expected = torch.tensor([[[ [1., 1., 1., 0.], [1., 1., 1., 0.], [0., 0., 0., 0.], ]]], device=device, dtype=dtype).repeat(2, 1, 1, 1) grid = kornia.utils.create_meshgrid(height, width, normalized_coordinates=False, device=device).to(dtype) grid += 1. # apply shift in both x/y direction input_warped = kornia.remap(inp, grid[..., 0], grid[..., 1], align_corners=True) assert_allclose(input_warped, expected, rtol=1e-4, atol=1e-4)
def test_smoke(self, device): height, width = 3, 4 input = torch.ones(1, 1, height, width).to(device) grid = kornia.utils.create_meshgrid( height, width, normalized_coordinates=False).to(device) input_warped = kornia.remap(input, grid[..., 0], grid[..., 1]) assert_allclose(input, input_warped)
def test_normalized_coordinates(self, device, dtype): height, width = 3, 4 normalized_coordinates = True inp = torch.tensor([[[[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]]], device=device, dtype=dtype).repeat(2, 1, 1, 1) expected = torch.tensor([[[[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]]], device=device, dtype=dtype).repeat(2, 1, 1, 1) grid = kornia.utils.create_meshgrid( height, width, normalized_coordinates=normalized_coordinates, device=device).to(dtype) # Normalized input coordinates input_warped = kornia.remap( inp, grid[..., 0], grid[..., 1], align_corners=True, normalized_coordinates=normalized_coordinates) assert_close(input_warped, expected, rtol=1e-4, atol=1e-4)
def test_smoke(self, device, dtype): height, width = 3, 4 input = torch.ones(1, 1, height, width, device=device, dtype=dtype) grid = kornia.utils.create_meshgrid( height, width, normalized_coordinates=False, device=device).to(dtype) input_warped = kornia.remap(input, grid[..., 0], grid[..., 1], align_corners=True) assert_allclose(input, input_warped, rtol=1e-4, atol=1e-4)
def test_shift(self, device): height, width = 3, 4 inp = torch.tensor([[[ [1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.], ]]]).to(device) expected = torch.tensor([[[ [1., 1., 1., 0.], [1., 1., 1., 0.], [0., 0., 0., 0.], ]]]).to(device) grid = kornia.utils.create_meshgrid( height, width, normalized_coordinates=False).to(device) grid += 1. # apply shift in both x/y direction input_warped = kornia.remap(inp, grid[..., 0], grid[..., 1]) assert_allclose(input_warped, expected)
def op_script(input, map1, map2): return kornia.remap(input, map1, map2)