def test_exception(self, device, dtype):
with pytest.raises(ValueError):
im = torch.rand(5, 5, device=device, dtype=dtype)
K = torch.rand(3, 3, device=device, dtype=dtype)
distCoeff = torch.rand(4, device=device, dtype=dtype)
undistort_image(im, K, distCoeff)
with pytest.raises(ValueError):
im = torch.rand(3, 5, 5, device=device, dtype=dtype)
K = torch.rand(4, 4, device=device, dtype=dtype)
distCoeff = torch.rand(4, device=device, dtype=dtype)
undistort_image(im, K, distCoeff)
with pytest.raises(ValueError):
im = torch.rand(3, 5, 5, device=device, dtype=dtype)
K = torch.rand(3, 3, device=device, dtype=dtype)
distCoeff = torch.rand(6, device=device, dtype=dtype)
undistort_image(im, K, distCoeff)
with pytest.raises(ValueError):
im = torch.randint(0,
256, (3, 5, 5),
device=device,
dtype=torch.uint8)
K = torch.rand(3, 3, device=device, dtype=dtype)
distCoeff = torch.rand(4, device=device, dtype=dtype)
undistort_image(im, K, distCoeff)
with pytest.raises(ValueError):
im = torch.rand(1, 1, 3, 5, 5, device=device, dtype=dtype)
K = torch.rand(1, 3, 3, device=device, dtype=dtype)
distCoeff = torch.rand(1, 4, device=device, dtype=dtype)
undistort_image(im, K, distCoeff)
def test_opencv(self, device, dtype):
im = torch.tensor(
[[[
[116, 75, 230, 5, 32],
[9, 182, 97, 213, 3],
[91, 10, 33, 141, 230],
[229, 63, 221, 244, 61],
[19, 137, 23, 59, 227],
]]],
device=device,
dtype=dtype,
)
K = torch.tensor([[2, 0, 2], [0, 2, 2], [0, 0, 1]],
device=device,
dtype=dtype)
dist = torch.tensor([0.2290, 0.9565, 0.0083, 0.0475],
device=device,
dtype=dtype)
# Expected output generated with OpenCV:
# import cv2
# imu_expected = cv2.undistort(np.uint8(im[0,0].numpy()), K.numpy(), dist.numpy())
imu_expected = torch.tensor(
[[[[0, 0, 0, 0, 0], [0, 124, 112, 82, 0], [0, 13, 33, 158, 0],
[0, 108, 197, 150, 0], [0, 0, 0, 0, 0]]]],
device=device,
dtype=dtype,
)
imu = undistort_image(im / 255.0, K, dist)
assert_close(imu, imu_expected / 255.0, rtol=1e-2, atol=1e-2)
def test_shape_minimum_dims(self, device, dtype):
im = torch.rand(3, 5, 5, device=device, dtype=dtype)
K = torch.rand(3, 3, device=device, dtype=dtype)
distCoeff = torch.rand(4, device=device, dtype=dtype)
imu = undistort_image(im, K, distCoeff)
assert imu.shape == (3, 5, 5)
def test_shape_extra_dims(self, device, dtype):
im = torch.rand(1, 1, 3, 5, 5, device=device,
dtype=dtype).tile(3, 2, 1, 1, 1)
K = torch.rand(1, 1, 3, 3, device=device, dtype=dtype).tile(3, 2, 1, 1)
distCoeff = torch.rand(1, 1, 4, device=device,
dtype=dtype).tile(3, 2, 1)
imu = undistort_image(im, K, distCoeff)
assert imu.shape == (3, 2, 3, 5, 5)
torch.testing.assert_allclose(imu[0], imu[1])