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)
