File size: 6,139 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import pytest
import torch
from torch.autograd import gradcheck
import kornia
import kornia.testing as utils # test utils
from kornia.testing import assert_close
class TestPyrUp:
def test_shape(self, device, dtype):
inp = torch.zeros(1, 2, 4, 4, device=device, dtype=dtype)
pyr = kornia.geometry.PyrUp()
assert pyr(inp).shape == (1, 2, 8, 8)
def test_shape_batch(self, device, dtype):
inp = torch.zeros(2, 2, 4, 4, device=device, dtype=dtype)
pyr = kornia.geometry.PyrUp()
assert pyr(inp).shape == (2, 2, 8, 8)
def test_gradcheck(self, device, dtype):
img = torch.rand(1, 2, 5, 4, device=device, dtype=dtype)
img = utils.tensor_to_gradcheck_var(img) # to var
assert gradcheck(kornia.geometry.pyrup, (img,), raise_exception=True)
def test_jit(self, device, dtype):
img = torch.rand(2, 3, 4, 5, device=device, dtype=dtype)
op = kornia.geometry.pyrup
op_jit = torch.jit.script(op)
assert_close(op(img), op_jit(img))
class TestPyrDown:
def test_shape(self, device, dtype):
inp = torch.zeros(1, 2, 4, 4, device=device, dtype=dtype)
pyr = kornia.geometry.PyrDown()
assert pyr(inp).shape == (1, 2, 2, 2)
def test_shape_batch(self, device, dtype):
inp = torch.zeros(2, 2, 4, 4, device=device, dtype=dtype)
pyr = kornia.geometry.PyrDown()
assert pyr(inp).shape == (2, 2, 2, 2)
def test_symmetry_preserving(self, device, dtype):
inp = torch.zeros(1, 1, 6, 6, device=device, dtype=dtype)
inp[:, :, 2:4, 2:4] = 1.0
pyr_out = kornia.geometry.PyrDown()(inp).squeeze()
assert_close(pyr_out, pyr_out.flip(0))
assert_close(pyr_out, pyr_out.flip(1))
def test_gradcheck(self, device, dtype):
img = torch.rand(1, 2, 5, 4, device=device, dtype=dtype)
img = utils.tensor_to_gradcheck_var(img) # to var
assert gradcheck(kornia.geometry.pyrdown, (img,), raise_exception=True)
def test_jit(self, device, dtype):
img = torch.rand(2, 3, 4, 5, device=device, dtype=dtype)
op = kornia.geometry.pyrdown
op_jit = torch.jit.script(op)
assert_close(op(img), op_jit(img))
class TestScalePyramid:
def test_shape_tuple(self, device, dtype):
inp = torch.zeros(3, 2, 41, 41, device=device, dtype=dtype)
SP = kornia.geometry.ScalePyramid(n_levels=1, min_size=30)
out = SP(inp)
assert len(out) == 3
assert len(out[0]) == 1
assert len(out[1]) == 1
assert len(out[2]) == 1
def test_shape_batch(self, device, dtype):
inp = torch.zeros(3, 2, 31, 31, device=device, dtype=dtype)
SP = kornia.geometry.ScalePyramid(n_levels=1)
sp, _, _ = SP(inp)
assert sp[0].shape == (3, 2, 3 + 1, 31, 31)
def test_shape_batch_double(self, device, dtype):
inp = torch.zeros(3, 2, 31, 31, device=device, dtype=dtype)
SP = kornia.geometry.ScalePyramid(n_levels=1, double_image=True)
sp, _, _ = SP(inp)
assert sp[0].shape == (3, 2, 1 + 3, 62, 62)
def test_n_levels_shape(self, device, dtype):
inp = torch.zeros(1, 1, 32, 32, device=device, dtype=dtype)
SP = kornia.geometry.ScalePyramid(n_levels=3)
sp, _, _ = SP(inp)
assert sp[0].shape == (1, 1, 3 + 3, 32, 32)
def test_blur_order(self, device, dtype):
inp = torch.rand(1, 1, 31, 31, device=device, dtype=dtype)
SP = kornia.geometry.ScalePyramid(n_levels=3)
sp, _, _ = SP(inp)
for _, pyr_level in enumerate(sp):
for _, img in enumerate(pyr_level):
img = img.squeeze().view(3, -1)
max_per_blur_level_val, _ = img.max(dim=1)
assert torch.argmax(max_per_blur_level_val).item() == 0
def test_symmetry_preserving(self, device, dtype):
PS = 16
R = 2
inp = torch.zeros(1, 1, PS, PS, device=device, dtype=dtype)
inp[..., PS // 2 - R : PS // 2 + R, PS // 2 - R : PS // 2 + R] = 1.0
SP = kornia.geometry.ScalePyramid(n_levels=3)
sp, _, _ = SP(inp)
for _, pyr_level in enumerate(sp):
for _, img in enumerate(pyr_level):
img = img.squeeze()
assert_close(img, img.flip(1))
assert_close(img, img.flip(2))
def test_gradcheck(self, device, dtype):
img = torch.rand(1, 2, 7, 9, device=device, dtype=dtype)
img = utils.tensor_to_gradcheck_var(img) # to var
from kornia.geometry import ScalePyramid as SP
def sp_tuple(img):
sp, _, _ = SP()(img)
return tuple(sp)
assert gradcheck(sp_tuple, (img,), raise_exception=True, nondet_tol=1e-4)
class TestBuildPyramid:
def test_smoke(self, device, dtype):
input = torch.ones(1, 2, 4, 5, device=device, dtype=dtype)
pyramid = kornia.geometry.transform.build_pyramid(input, max_level=1)
assert len(pyramid) == 1
assert pyramid[0].shape == (1, 2, 4, 5)
@pytest.mark.parametrize("batch_size", (1, 2, 3))
@pytest.mark.parametrize("channels", (1, 3))
@pytest.mark.parametrize("max_level", (2, 3, 4))
def test_num_levels(self, batch_size, channels, max_level, device, dtype):
height, width = 16, 20
input = torch.rand(batch_size, channels, height, width, device=device, dtype=dtype)
pyramid = kornia.geometry.transform.build_pyramid(input, max_level)
assert len(pyramid) == max_level
for i in range(1, max_level):
img = pyramid[i]
denom = 2 ** i
expected_shape = (batch_size, channels, height // denom, width // denom)
assert img.shape == expected_shape
def test_gradcheck(self, device, dtype):
max_level = 1
batch_size, channels, height, width = 1, 2, 7, 9
img = torch.rand(batch_size, channels, height, width, device=device, dtype=dtype)
img = utils.tensor_to_gradcheck_var(img) # to var
assert gradcheck(kornia.geometry.transform.build_pyramid, (img, max_level), raise_exception=True)
|