|
|
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)) |
|
|
|
