File size: 2,181 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 | import pytest
import torch
from torch.autograd import gradcheck
from kornia.geometry.calibration.distort import distort_points
from kornia.testing import assert_close
class TestDistortPoints:
def test_smoke(self, device, dtype):
points = torch.rand(1, 2, device=device, dtype=dtype)
K = torch.rand(3, 3, device=device, dtype=dtype)
distCoeff = torch.rand(4, device=device, dtype=dtype)
pointsu = distort_points(points, K, distCoeff)
assert points.shape == pointsu.shape
def test_smoke_batch(self, device, dtype):
points = torch.rand(1, 1, 2, device=device, dtype=dtype)
K = torch.rand(1, 3, 3, device=device, dtype=dtype)
distCoeff = torch.rand(1, 4, device=device, dtype=dtype)
pointsu = distort_points(points, K, distCoeff)
assert points.shape == pointsu.shape
@pytest.mark.parametrize(
"batch_size, num_points, num_distcoeff", [(1, 3, 4), (2, 4, 5), (3, 5, 8), (4, 6, 12), (5, 7, 14)]
)
def test_shape(self, batch_size, num_points, num_distcoeff, device, dtype):
B, N, Ndist = batch_size, num_points, num_distcoeff
points = torch.rand(B, N, 2, device=device, dtype=dtype)
K = torch.rand(B, 3, 3, device=device, dtype=dtype)
distCoeff = torch.rand(B, Ndist, device=device, dtype=dtype)
pointsu = distort_points(points, K, distCoeff)
assert pointsu.shape == (B, N, 2)
def test_gradcheck(self, device):
points = torch.rand(1, 8, 2, device=device, dtype=torch.float64, requires_grad=True)
K = torch.rand(1, 3, 3, device=device, dtype=torch.float64)
distCoeff = torch.rand(1, 4, device=device, dtype=torch.float64)
assert gradcheck(distort_points, (points, K, distCoeff), raise_exception=True)
def test_jit(self, device, dtype):
points = torch.rand(1, 1, 2, device=device, dtype=dtype)
K = torch.rand(1, 3, 3, device=device, dtype=dtype)
distCoeff = torch.rand(1, 4, device=device, dtype=dtype)
inputs = (points, K, distCoeff)
op = distort_points
op_jit = torch.jit.script(op)
assert_close(op(*inputs), op_jit(*inputs))
|