def test_jit(self, device, dtype):
shear = torch.tensor([[0.5, 0.0]], device=device, dtype=dtype)
batch_size, channels, height, width = 2, 3, 64, 64
img = torch.ones(batch_size, channels, height, width, device=device, dtype=dtype)
trans = kornia.Shear(shear)
trans_traced = torch.jit.trace(kornia.Shear(shear), img)
assert_close(trans(img), trans_traced(img), atol=1e-4, rtol=1e-4)
def test_jit(self, device):
shear = torch.tensor([[0.5, 0.0]]).to(device)
batch_size, channels, height, width = 2, 3, 64, 64
img = torch.ones(batch_size, channels, height, width).to(device)
trans = kornia.Shear(shear)
trans_traced = torch.jit.trace(kornia.Shear(shear), img)
assert_allclose(trans(img), trans_traced(img))
def test_shear_batch2(self):
# prepare input data
inp = torch.tensor([[
[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]
]]).repeat(2, 1, 1, 1)
expected = torch.tensor([[[
[1., 1., 1., 1.],
[.5, 1., 1., 1.],
[0., 1., 1., 1.],
[0., .5, 1., 1.]
]], [[
[1., .5, 0., 0.],
[1., 1., 1., .5],
[1., 1., 1., 1.],
[1., 1., 1., 1.]
]]])
# prepare transformation
shear = torch.tensor([[0.5, 0.0], [0.0, 0.5]])
transform = kornia.Shear(shear)
assert_allclose(transform(inp), expected)
def test_shear_batch2(self, device, dtype):
# prepare input data
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],
[1.0, 1.0, 1.0, 1.0]]],
device=device,
dtype=dtype,
).repeat(2, 1, 1, 1)
expected = torch.tensor(
[
[[[0.75, 1.0, 1.0, 1.0], [0.25, 1.0, 1.0, 1.0],
[0.0, 0.75, 1.0, 1.0], [0.0, 0.25, 1.0, 1.0]]],
[[[0.75, 0.25, 0.0, 0.0], [1.0, 1.0, 0.75, 0.25],
[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]],
],
device=device,
dtype=dtype,
)
# prepare transformation
shear = torch.tensor([[0.5, 0.0], [0.0, 0.5]],
device=device,
dtype=dtype)
transform = kornia.Shear(shear)
assert_allclose(transform(inp), expected, atol=1e-4, rtol=1e-4)
def test_shear_y(self, device):
# prepare input data
inp = torch.tensor([[[1., 1., 1., 1.], [1., 1., 1., 1.],
[1., 1., 1., 1.], [1., 1., 1., 1.]]]).to(device)
expected = torch.tensor([[[0.75, 0.25, 0., 0.], [1., 1., 0.75, 0.25],
[1., 1., 1., 1.], [1., 1., 1.,
1.]]]).to(device)
# prepare transformation
shear = torch.tensor([[0.0, 0.5]]).to(device)
transform = kornia.Shear(shear)
assert_allclose(transform(inp), expected)
def test_shear_batch2_broadcast(self, device):
# prepare input data
inp = torch.tensor([[[1., 1., 1., 1.], [1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]]]).repeat(2, 1, 1, 1).to(device)
expected = torch.tensor([[[[0.75, 1., 1., 1.], [0.25, 1., 1., 1.],
[0., 0.75, 1., 1.], [0., 0.25, 1.,
1.]]]]).to(device)
# prepare transformation
shear = torch.tensor([[0.5, 0.0]]).to(device)
transform = kornia.Shear(shear)
assert_allclose(transform(inp), expected)
def test_shear_x(self, device, dtype):
# prepare input data
inp = torch.tensor([[[1., 1., 1., 1.], [1., 1., 1., 1.],
[1., 1., 1., 1.], [1., 1., 1., 1.]]],
device=device,
dtype=dtype)
expected = torch.tensor([[[0.75, 1., 1., 1.], [0.25, 1., 1., 1.],
[0., 0.75, 1., 1.], [0., 0.25, 1., 1.]]],
device=device,
dtype=dtype)
# prepare transformation
shear = torch.tensor([[0.5, 0.0]], device=device, dtype=dtype)
transform = kornia.Shear(shear)
assert_allclose(transform(inp), expected, atol=1e-4, rtol=1e-4)
def test_shear_y(self, device, dtype):
# prepare input data
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], [1.0, 1.0, 1.0, 1.0]]],
device=device,
dtype=dtype,
)
expected = torch.tensor(
[[[0.75, 0.25, 0.0, 0.0], [1.0, 1.0, 0.75, 0.25], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]],
device=device,
dtype=dtype,
)
# prepare transformation
shear = torch.tensor([[0.0, 0.5]], device=device, dtype=dtype)
transform = kornia.Shear(shear, align_corners=False)
assert_close(transform(inp), expected, atol=1e-4, rtol=1e-4)
