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