class TestLoFTR:
def test_pretrained_outdoor_smoke(self, device, dtype):
loftr = LoFTR('outdoor').to(device, dtype)
assert loftr is not None
def test_pretrained_indoor_smoke(self, device, dtype):
loftr = LoFTR('indoor').to(device, dtype)
assert loftr is not None
@pytest.mark.skipif(torch_version_geq(1, 10), reason="RuntimeError: CUDA out of memory with pytorch>=1.10")
@pytest.mark.skipif(sys.platform == "win32",
reason="this test takes so much memory in the CI with Windows")
@pytest.mark.parametrize("data", ["loftr_fund"], indirect=True)
def test_pretrained_indoor(self, device, dtype, data):
loftr = LoFTR('indoor').to(device, dtype)
data_dev = utils.dict_to(data, device, dtype)
with torch.no_grad():
out = loftr(data_dev)
assert_close(out['keypoints0'], data_dev["loftr_indoor_tentatives0"])
assert_close(out['keypoints1'], data_dev["loftr_indoor_tentatives1"])
@pytest.mark.skipif(torch_version_geq(1, 10), reason="RuntimeError: CUDA out of memory with pytorch>=1.10")
@pytest.mark.skipif(sys.platform == "win32",
reason="this test takes so much memory in the CI with Windows")
@pytest.mark.parametrize("data", ["loftr_homo"], indirect=True)
def test_pretrained_outdoor(self, device, dtype, data):
loftr = LoFTR('outdoor').to(device, dtype)
data_dev = utils.dict_to(data, device, dtype)
with torch.no_grad():
out = loftr(data_dev)
assert_close(out['keypoints0'], data_dev["loftr_outdoor_tentatives0"])
assert_close(out['keypoints1'], data_dev["loftr_outdoor_tentatives1"])
@pytest.mark.skip("Takes too long time (but works)")
def test_gradcheck(self, device):
patches = torch.rand(1, 1, 32, 32, device=device)
patches05 = resize(patches, (48, 48))
patches = utils.tensor_to_gradcheck_var(patches) # to var
patches05 = utils.tensor_to_gradcheck_var(patches05) # to var
loftr = LoFTR().to(patches.device, patches.dtype)
def proxy_forward(x, y):
return loftr.forward({"image0": x, "image1": y})["keypoints0"]
assert gradcheck(proxy_forward, (patches, patches05), eps=1e-4, atol=1e-4, raise_exception=True)
@pytest.mark.skip("does not like transformer.py:L99, zip iteration")
def test_jit(self, device, dtype):
B, C, H, W = 1, 1, 32, 32
patches = torch.rand(B, C, H, W, device=device, dtype=dtype)
patches2x = resize(patches, (48, 48))
input = {"image0": patches, "image1": patches2x}
model = LoFTR().to(patches.device, patches.dtype).eval()
model_jit = torch.jit.script(model)
out = model(input)
out_jit = model_jit(input)
for k, v in out.items():
assert_close(v, out_jit[k])
def test_batch_random_affine_3d(self, device, dtype):
# TODO(jian): crashes with pytorch 1.10, cuda and fp64
if torch_version_geq(
1, 10) and "cuda" in str(device) and dtype == torch.float64:
pytest.skip(
"AssertionError: assert tensor(False, device='cuda:0')")
f = RandomAffine3D((0, 0, 0), p=1.0,
return_transform=True) # No rotation
tensor = torch.tensor(
[[[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]]]],
device=device,
dtype=dtype) # 1 x 1 x 1 x 3 x 3
expected = torch.tensor(
[[[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]]]],
device=device,
dtype=dtype) # 1 x 1 x 1 x 3 x 3
expected_transform = torch.tensor(
[[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0]]],
device=device,
dtype=dtype,
) # 1 x 4 x 4
tensor = tensor.repeat(5, 3, 1, 1, 1) # 5 x 3 x 3 x 3 x 3
expected = expected.repeat(5, 3, 1, 1, 1) # 5 x 3 x 3 x 3 x 3
expected_transform = expected_transform.repeat(5, 1, 1) # 5 x 4 x 4
assert (f(tensor)[0] == expected).all()
assert (f(tensor)[1] == expected_transform).all()
def safe_solve_with_mask(B: torch.Tensor, A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
r"""Helper function, which avoids crashing because of singular matrix input and outputs the
mask of valid solution"""
if not torch_version_geq(1, 10):
sol, lu = _torch_solve_cast(B, A)
warnings.warn('PyTorch version < 1.10, solve validness mask maybe not correct', RuntimeWarning)
return sol, lu, torch.ones(len(A), dtype=torch.bool, device=A.device)
# Based on https://github.com/pytorch/pytorch/issues/31546#issuecomment-694135622
if not isinstance(B, torch.Tensor):
raise AssertionError(f"B must be torch.Tensor. Got: {type(B)}.")
dtype: torch.dtype = B.dtype
if dtype not in (torch.float32, torch.float64):
dtype = torch.float32
A_LU, pivots, info = torch.lu(A.to(dtype), get_infos=True)
valid_mask: torch.Tensor = info == 0
X = torch.lu_solve(B.to(dtype), A_LU, pivots)
return X.to(B.dtype), A_LU.to(A.dtype), valid_mask
def safe_inverse_with_mask(A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
r"""Helper function, which avoids crashing because of non-invertable matrix input and outputs the
mask of valid solution"""
# Based on https://github.com/pytorch/pytorch/issues/31546#issuecomment-694135622
if not torch_version_geq(1, 9):
inv = _torch_inverse_cast(A)
warnings.warn('PyTorch version < 1.9, inverse validness mask maybe not correct', RuntimeWarning)
return inv, torch.ones(len(A), dtype=torch.bool, device=A.device)
if not isinstance(A, torch.Tensor):
raise AssertionError(f"A must be torch.Tensor. Got: {type(A)}.")
dtype_original: torch.dtype = A.dtype
if dtype_original not in (torch.float32, torch.float64):
dtype = torch.float32
else:
dtype = dtype_original
from torch.linalg import inv_ex # type: ignore # (not available in 1.8.1)
inverse, info = inv_ex(A.to(dtype))
mask = info == 0
return inverse.to(dtype_original), mask
class TestWarpImage:
@pytest.mark.parametrize('batch_size', [1, 3])
def test_smoke(self, batch_size, device, dtype):
src, dst = _sample_points(batch_size, device)
tensor = torch.rand(batch_size, 3, 32, 32, device=device)
kernel, affine = kornia.geometry.transform.get_tps_transform(src, dst)
warp = kornia.geometry.transform.warp_image_tps(
tensor, dst, kernel, affine)
assert warp.shape == tensor.shape
@pytest.mark.skipif(
torch_version_geq(1, 10),
reason=
"for some reason the solver detects singular matrices in pytorch >=1.10."
)
@pytest.mark.parametrize('batch_size', [1, 3])
def test_warp(self, batch_size, device, dtype):
src = torch.tensor([[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0],
[1.0, -1.0], [0.0, 0.0]]],
device=device).repeat(batch_size, 1, 1)
# zoom in by a factor of 2
dst = src.clone() * 2.0
tensor = torch.zeros(batch_size, 3, 8, 8, device=device)
tensor[:, :, 2:6, 2:6] = 1.0
expected = torch.ones_like(tensor)
# nn.grid_sample interpolates the at the edges it seems, so the boundaries have values < 1
expected[:, :, [0, -1], :] *= 0.5
expected[:, :, :, [0, -1]] *= 0.5
kernel, affine = kornia.geometry.transform.get_tps_transform(dst, src)
warp = kornia.geometry.transform.warp_image_tps(
tensor, src, kernel, affine)
assert_close(warp, expected, atol=1e-4, rtol=1e-4)
@pytest.mark.parametrize('batch_size', [1, 3])
def test_exception(self, batch_size, device, dtype):
image = torch.rand(batch_size, 3, 32, 32)
dst = torch.rand(batch_size, 5, 2)
kernel = torch.zeros_like(dst)
affine = torch.zeros(batch_size, 3, 2)
with pytest.raises(TypeError):
assert kornia.geometry.transform.warp_image_tps(
image.numpy(), dst, kernel, affine)
with pytest.raises(TypeError):
assert kornia.geometry.transform.warp_image_tps(
image, dst.numpy(), kernel, affine)
with pytest.raises(TypeError):
assert kornia.geometry.transform.warp_image_tps(
image, dst, kernel.numpy(), affine)
with pytest.raises(TypeError):
assert kornia.geometry.transform.warp_image_tps(
image, dst, kernel, affine.numpy())
with pytest.raises(ValueError):
image_bad = torch.rand(batch_size, 32, 32)
assert kornia.geometry.transform.warp_image_tps(
image_bad, dst, kernel, affine)
with pytest.raises(ValueError):
dst_bad = torch.rand(batch_size, 5)
assert kornia.geometry.transform.warp_image_tps(
image, dst_bad, kernel, affine)
with pytest.raises(ValueError):
kernel_bad = torch.rand(batch_size, 5)
assert kornia.geometry.transform.warp_image_tps(
image, dst, kernel_bad, affine)
with pytest.raises(ValueError):
affine_bad = torch.rand(batch_size, 3)
assert kornia.geometry.transform.warp_image_tps(
image, dst, kernel, affine_bad)
@pytest.mark.grad
@pytest.mark.parametrize('batch_size', [1, 3])
def test_gradcheck(self, batch_size, device, dtype):
opts = dict(device=device, dtype=torch.float64)
src, dst = _sample_points(batch_size, **opts)
kernel, affine = kornia.geometry.transform.get_tps_transform(src, dst)
image = torch.rand(batch_size, 3, 8, 8, requires_grad=True, **opts)
assert gradcheck(
kornia.geometry.transform.warp_image_tps,
(image, dst, kernel, affine),
raise_exception=True,
atol=1e-4,
rtol=1e-4,
)
@pytest.mark.jit
@pytest.mark.parametrize('batch_size', [1, 3])
def test_jit(self, batch_size, device, dtype):
src, dst = _sample_points(batch_size, device)
kernel, affine = kornia.geometry.transform.get_tps_transform(src, dst)
image = torch.rand(batch_size, 3, 32, 32, device=device)
op = kornia.geometry.transform.warp_image_tps
op_jit = torch.jit.script(op)
assert_close(op(image, dst, kernel, affine),
op_jit(image, dst, kernel, affine),
rtol=1e-4,
atol=1e-4)
def conv_quad_interp3d(input: torch.Tensor,
strict_maxima_bonus: float = 10.0,
eps: float = 1e-7) -> Tuple[torch.Tensor, torch.Tensor]:
r"""Compute the single iteration of quadratic interpolation of the extremum (max or min).
Args:
input: the given heatmap with shape :math:`(N, C, D_{in}, H_{in}, W_{in})`.
strict_maxima_bonus: pixels, which are strict maxima will score (1 + strict_maxima_bonus) * value.
This is needed for mimic behavior of strict NMS in classic local features
eps: parameter to control the hessian matrix ill-condition number.
Returns:
the location and value per each 3x3x3 window which contains strict extremum, similar to one done is SIFT.
:math:`(N, C, 3, D_{out}, H_{out}, W_{out})`, :math:`(N, C, D_{out}, H_{out}, W_{out})`,
where
.. math::
D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] -
(\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor
.. math::
H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] -
(\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor
.. math::
W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] -
(\text{kernel\_size}[2] - 1) - 1}{\text{stride}[2]} + 1\right\rfloor
Examples:
>>> input = torch.randn(20, 16, 3, 50, 32)
>>> nms_coords, nms_val = conv_quad_interp3d(input, 1.0)
"""
if not torch.is_tensor(input):
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
if not len(input.shape) == 5:
raise ValueError(
f"Invalid input shape, we expect BxCxDxHxW. Got: {input.shape}")
B, CH, D, H, W = input.shape
grid_global: torch.Tensor = create_meshgrid3d(D,
H,
W,
False,
device=input.device).permute(
0, 4, 1, 2, 3)
grid_global = grid_global.to(input.dtype)
# to determine the location we are solving system of linear equations Ax = b, where b is 1st order gradient
# and A is Hessian matrix
b: torch.Tensor = spatial_gradient3d(input, order=1, mode='diff') #
b = b.permute(0, 1, 3, 4, 5, 2).reshape(-1, 3, 1)
A: torch.Tensor = spatial_gradient3d(input, order=2, mode='diff')
A = A.permute(0, 1, 3, 4, 5, 2).reshape(-1, 6)
dxx = A[..., 0]
dyy = A[..., 1]
dss = A[..., 2]
dxy = 0.25 * A[..., 3] # normalization to match OpenCV implementation
dys = 0.25 * A[..., 4] # normalization to match OpenCV implementation
dxs = 0.25 * A[..., 5] # normalization to match OpenCV implementation
Hes = torch.stack([dxx, dxy, dxs, dxy, dyy, dys, dxs, dys, dss],
dim=-1).view(-1, 3, 3)
if not torch_version_geq(1, 10):
# The following is needed to avoid singular cases
Hes += torch.rand(Hes[0].size(), device=Hes.device).abs()[None] * eps
nms_mask: torch.Tensor = nms3d(input, (3, 3, 3), True)
x_solved: torch.Tensor = torch.zeros_like(b)
x_solved_masked, _, solved_correctly = safe_solve_with_mask(
b[nms_mask.view(-1)], Hes[nms_mask.view(-1)])
# Kill those points, where we cannot solve
new_nms_mask = nms_mask.masked_scatter(nms_mask, solved_correctly)
x_solved.masked_scatter_(new_nms_mask.view(-1, 1, 1),
x_solved_masked[solved_correctly])
dx: torch.Tensor = -x_solved
# Ignore ones, which are far from window center
mask1 = dx.abs().max(dim=1, keepdim=True)[0] > 0.7
dx.masked_fill_(mask1.expand_as(dx), 0)
dy: torch.Tensor = 0.5 * torch.bmm(b.permute(0, 2, 1), dx)
y_max = input + dy.view(B, CH, D, H, W)
if strict_maxima_bonus > 0:
y_max += strict_maxima_bonus * new_nms_mask.to(input.dtype)
dx_res: torch.Tensor = dx.flip(1).reshape(B, CH, D, H, W,
3).permute(0, 1, 5, 2, 3, 4)
coords_max: torch.Tensor = grid_global.repeat(B, 1, 1, 1, 1).unsqueeze(1)
coords_max = coords_max + dx_res
return coords_max, y_max
