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