Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
@pytest.mark.skip(reason="turn off all jit for a while")
def test_jit(self, device):
op = kornia.denormalize_pixel_coordinates
op_script = torch.jit.script(op)
height, width = 3, 4
grid = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=True).to(device)
actual = op_script(grid, height, width)
expected = op(grid, height, width)
assert_allclose(actual, expected)
def test_normalize_pixel_grid():
# generate input data
batch_size = 1
height, width = 2, 4
# create points grid
grid_norm = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=True)
grid_norm = torch.unsqueeze(grid_norm, dim=0)
grid_pix = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=False)
grid_pix = torch.unsqueeze(grid_pix, dim=0)
# grid from pixel space to normalized
norm_trans_pix = kornia.normal_transform_pixel(height, width) # 1x3x3
pix_trans_norm = torch.inverse(norm_trans_pix) # 1x3x3
# transform grids
grid_pix_to_norm = kornia.transform_points(norm_trans_pix, grid_pix)
grid_norm_to_pix = kornia.transform_points(pix_trans_norm, grid_norm)
assert_allclose(grid_pix, grid_norm_to_pix)
assert_allclose(grid_norm, grid_pix_to_norm)
def test_tensor_bhw2(self, device):
height, width = 3, 4
grid = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=True).to(device)
expected = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=False).to(device)
grid_norm = kornia.denormalize_pixel_coordinates(
grid, height, width)
assert_allclose(grid_norm, expected)
def test_shift_batch_broadcast(self):
height, width = 3, 4
inp = torch.tensor([[[
[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.],
]]]).repeat(2, 1, 1, 1)
expected = torch.tensor([[[
[1., 1., 1., 0.],
[1., 1., 1., 0.],
[0., 0., 0., 0.],
]]])
grid = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=False)
grid += 1. # apply shift in both x/y direction
input_warped = kornia.remap(inp, grid[..., 0], grid[..., 1])
assert_allclose(input_warped, expected)
"""
if not torch.is_tensor(input):
raise TypeError("Input input type is not a torch.Tensor. Got {}"
.format(type(input)))
if not len(input.shape) == 4:
raise ValueError("Invalid input shape, we expect BxCxHxW. Got: {}"
.format(input.shape))
# unpack shapes and create view from input tensor
batch_size, channels, height, width = input.shape
x: torch.Tensor = input.view(batch_size, channels, -1)
# compute softmax along the feature map
x_soft: torch.Tensor = F.softmax(x * temperature, dim=-1)
# create coordinates grid
grid: torch.Tensor = create_meshgrid(
height, width, normalized_coordinates)
grid = grid.to(input.device).to(input.dtype)
pos_x: torch.Tensor = grid[..., 0].reshape(-1)
pos_y: torch.Tensor = grid[..., 1].reshape(-1)
# compute the expected coordinates
expected_y: torch.Tensor = torch.sum(pos_y * x_soft, dim=-1, keepdim=True)
expected_x: torch.Tensor = torch.sum(pos_x * x_soft, dim=-1, keepdim=True)
output: torch.Tensor = torch.cat([expected_x, expected_y], dim=-1)
return output.view(batch_size, channels, 2) # BxNx2
raise TypeError(f"Input depht type is not a torch.Tensor. Got {type(depth)}.")
if not len(depth.shape) == 4 and depth.shape[-3] == 1:
raise ValueError(f"Input depth musth have a shape (B, 1, H, W). Got: {depth.shape}")
if not isinstance(camera_matrix, torch.Tensor):
raise TypeError(f"Input camera_matrix type is not a torch.Tensor. "
f"Got {type(camera_matrix)}.")
if not len(camera_matrix.shape) == 3 and camera_matrix.shape[-2:] == (3, 3):
raise ValueError(f"Input camera_matrix must have a shape (B, 3, 3). "
f"Got: {camera_matrix.shape}.")
# create base coordinates grid
batch_size, _, height, width = depth.shape
points_2d: torch.Tensor = create_meshgrid(
height, width, normalized_coordinates=False) # 1xHxWx2
points_2d = points_2d.to(depth.device).to(depth.dtype)
# depth should come in Bx1xHxW
points_depth: torch.Tensor = depth.permute(0, 2, 3, 1) # 1xHxWx1
# project pixels to camera frame
camera_matrix_tmp: torch.Tensor = camera_matrix[:, None, None] # Bx1x1x3x3
points_3d: torch.Tensor = unproject_points(
points_2d, points_depth, camera_matrix_tmp, normalize=True) # BxHxWx3
return points_3d.permute(0, 3, 1, 2) # Bx3xHxW
def _create_meshgrid(height: int, width: int) -> torch.Tensor:
grid: torch.Tensor = create_meshgrid(
height, width, normalized_coordinates=False) # 1xHxWx2
return convert_points_to_homogeneous(grid) # append ones to last dim
def __init__(
self,
height: int,
width: int,
mode: str = 'bilinear',
padding_mode: str = 'zeros',
normalized_coordinates: bool = True) -> None:
super(HomographyWarper, self).__init__()
self.width: int = width
self.height: int = height
self.mode: str = mode
self.padding_mode: str = padding_mode
self.normalized_coordinates: bool = normalized_coordinates
# create base grid to compute the flow
self.grid: torch.Tensor = create_meshgrid(
height, width, normalized_coordinates=normalized_coordinates)