File size: 3,670 Bytes
36c95ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import pytest
import torch
from torch.autograd import gradcheck
import kornia
from kornia.geometry.transform import elastic_transform2d
from kornia.testing import assert_close
class TestElasticTransform:
def test_smoke(self, device, dtype):
image = torch.rand(1, 4, 5, 5, device=device, dtype=dtype)
noise = torch.rand(1, 2, 5, 5, device=device, dtype=dtype)
assert elastic_transform2d(image, noise) is not None
@pytest.mark.parametrize("batch, channels, height, width", [(1, 3, 3, 4), (2, 2, 2, 4), (1, 5, 4, 1)])
def test_cardinality(self, device, dtype, batch, channels, height, width):
shape = batch, channels, height, width
img = torch.ones(shape, device=device, dtype=dtype)
noise = torch.ones((batch, 2, height, width), device=device, dtype=dtype)
assert elastic_transform2d(img, noise).shape == shape
def test_exception(self, device, dtype):
with pytest.raises(TypeError):
assert elastic_transform2d([0.0])
with pytest.raises(TypeError):
assert elastic_transform2d(torch.tensor(), 1)
with pytest.raises(ValueError):
img = torch.ones(1, 1, 1, device=device, dtype=dtype)
noise = torch.ones(1, 2, 1, 1, device=device, dtype=dtype)
assert elastic_transform2d(img, noise)
with pytest.raises(ValueError):
img = torch.ones(1, 1, 1, 1, device=device, dtype=dtype)
noise = torch.ones(1, 3, 1, 1, device=device, dtype=dtype)
assert elastic_transform2d(img, noise)
@pytest.mark.parametrize(
"kernel_size, sigma, alpha",
[
[(3, 3), (4.0, 4.0), (32.0, 32.0)],
[(5, 3), (4.0, 8.0), (16.0, 32.0)],
[(5, 5), torch.tensor([2.0, 8.0]), torch.tensor([16.0, 64.0])],
],
)
def test_valid_paramters(self, device, dtype, kernel_size, sigma, alpha):
image = torch.rand(1, 4, 5, 5, device=device, dtype=dtype)
noise = torch.rand(1, 2, 5, 5, device=device, dtype=dtype)
if isinstance(sigma, torch.Tensor):
sigma = sigma.to(device, dtype)
if isinstance(alpha, torch.Tensor):
alpha = alpha.to(device, dtype)
assert elastic_transform2d(image, noise, kernel_size, sigma, alpha) is not None
def test_values(self, device, dtype):
image = torch.tensor(
[[[[0.0018, 0.7521, 0.7550], [0.2053, 0.4249, 0.1369], [0.1027, 0.3992, 0.8773]]]],
device=device,
dtype=dtype,
)
noise = torch.ones(1, 2, 3, 3, device=device, dtype=dtype)
expected = torch.tensor(
[[[[0.0005, 0.3795, 0.1905], [0.1034, 0.4235, 0.0702], [0.0259, 0.2007, 0.2193]]]],
device=device,
dtype=dtype,
)
actual = elastic_transform2d(image, noise)
assert_close(actual, expected, atol=1e-3, rtol=1e-3)
@pytest.mark.parametrize("requires_grad", [True, False])
def test_gradcheck(self, device, dtype, requires_grad):
image = torch.rand(1, 1, 3, 3, device=device, dtype=torch.float64, requires_grad=requires_grad)
noise = torch.rand(1, 2, 3, 3, device=device, dtype=torch.float64, requires_grad=not requires_grad)
assert gradcheck(elastic_transform2d, (image, noise), raise_exception=True)
def test_jit(self, device, dtype):
image = torch.rand(1, 4, 5, 5, device=device, dtype=dtype)
noise = torch.rand(1, 2, 5, 5, device=device, dtype=dtype)
op = kornia.geometry.transform.elastic_transform2d
op_jit = torch.jit.script(op)
assert_close(op(image, noise), op_jit(image, noise))
|