def test_jit_trace(self):
@torch.jit.script
def op_script(input, height, width):
return kornia.normalize_pixel_coordinates(input, height, width)
# 1. Trace op
height, width = 3, 4
grid = kornia.utils.create_meshgrid(height,
width,
normalized_coordinates=False)
op_traced = torch.jit.trace(op_script, (
grid,
torch.tensor(height),
torch.tensor(width),
))
# 2. Generate new input
height, width = 2, 5
grid = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=False).repeat(2, 1, 1, 1)
# 3. Evaluate
actual = op_traced(grid, torch.tensor(height), torch.tensor(width))
expected = kornia.normalize_pixel_coordinates(grid, height, width)
assert_allclose(actual, expected)
def test_tensor_bhw2(self, device):
height, width = 3, 4
grid = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=False).to(device)
expected = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=True).to(device)
grid_norm = kornia.normalize_pixel_coordinates(grid, height, width)
assert_allclose(grid_norm, expected)
def warp_frame_depth(
image_src: torch.Tensor,
depth_dst: torch.Tensor,
src_trans_dst: torch.Tensor,
camera_matrix: torch.Tensor,
normalize_points: bool = False,
sampling_mode='bilinear') -> torch.Tensor:
# TAKEN FROM KORNIA LIBRARY
if not isinstance(image_src, torch.Tensor):
raise TypeError(f"Input image_src type is not a torch.Tensor. Got {type(image_src)}.")
if not len(image_src.shape) == 4:
raise ValueError(f"Input image_src musth have a shape (B, D, H, W). Got: {image_src.shape}")
if not isinstance(depth_dst, torch.Tensor):
raise TypeError(f"Input depht_dst type is not a torch.Tensor. Got {type(depth_dst)}.")
if not len(depth_dst.shape) == 4 and depth_dst.shape[-3] == 1:
raise ValueError(f"Input depth_dst musth have a shape (B, 1, H, W). Got: {depth_dst.shape}")
if not isinstance(src_trans_dst, torch.Tensor):
raise TypeError(f"Input src_trans_dst type is not a torch.Tensor. "
f"Got {type(src_trans_dst)}.")
if not len(src_trans_dst.shape) == 3 and src_trans_dst.shape[-2:] == (3, 3):
raise ValueError(f"Input src_trans_dst must have a shape (B, 3, 3). "
f"Got: {src_trans_dst.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}.")
# unproject source points to camera frame
points_3d_dst: torch.Tensor = kornia.depth_to_3d(depth_dst, camera_matrix, normalize_points) # Bx3xHxW
# transform points from source to destination
points_3d_dst = points_3d_dst.permute(0, 2, 3, 1) # BxHxWx3
# apply transformation to the 3d points
points_3d_src = kornia.transform_points(src_trans_dst[:, None], points_3d_dst) # BxHxWx3
points_3d_src[:, :, :, 2] = torch.relu(points_3d_src[:, :, :, 2])
# project back to pixels
camera_matrix_tmp: torch.Tensor = camera_matrix[:, None, None] # Bx1x1xHxW
points_2d_src: torch.Tensor = kornia.project_points(points_3d_src, camera_matrix_tmp) # BxHxWx2
# normalize points between [-1 / 1]
height, width = depth_dst.shape[-2:]
points_2d_src_norm: torch.Tensor = kornia.normalize_pixel_coordinates(points_2d_src, height, width) # BxHxWx2
return torch.nn.functional.grid_sample(image_src, points_2d_src_norm, align_corners=True, mode=sampling_mode)
def test_list(self, device):
height, width = 3, 4
grid = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=False).to(device)
grid = grid.contiguous().view(-1, 2)
expected = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=True).to(device)
expected = expected.contiguous().view(-1, 2)
grid_norm = kornia.normalize_pixel_coordinates(grid, height, width)
assert_allclose(grid_norm, expected)
def test_jit(self):
@torch.jit.script
def op_script(input: torch.Tensor, height: int,
width: int) -> torch.Tensor:
return kornia.normalize_pixel_coordinates(input, height, width)
height, width = 3, 4
grid = kornia.utils.create_meshgrid(
height, width, normalized_coordinates=False)
actual = op_script(grid, height, width)
expected = kornia.normalize_pixel_coordinates(
grid, height, width)
assert_allclose(actual, expected)
def op_script(input, height, width):
return kornia.normalize_pixel_coordinates(input, height, width)
def op_script(input: torch.Tensor, height: int,
width: int) -> torch.Tensor:
return kornia.normalize_pixel_coordinates(input, height, width)
