def homography_warp(patch_src,
src_homo_dst,
dsize,
mode='bilinear',
padding_mode='zeros',
align_corners=False,
normalized_coordinates=True):
"""
Warp image patchs or tensors by normalized 2D homographies. See `HomographyWarper` for details.
"""
if not src_homo_dst.device == patch_src.device:
raise TypeError("Patch and homography must be on the same device. \
Got patch.device: {} src_H_dst.device: {}.".format(
patch_src.device, src_homo_dst.device))
height, width = dsize
grid = create_meshgrid(height,
width,
normalized_coordinates=normalized_coordinates)
warped_grid = warp_grid(grid, src_homo_dst)
return F.grid_sample(patch_src,
warped_grid,
mode=mode,
padding_mode=padding_mode,
align_corners=align_corners)
def depth_to_3d(depth, camera_matrix, normalize_points=False):
"""
Compute a 3d point per pixel given its depth value and the camera intrinsics.
"""
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 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 = 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 = depth.permute(0, 2, 3, 1) # 1xHxWx1
# project pixels to camera frame
camera_matrix_tmp = camera_matrix[:, None, None] # Bx1x1x3x3
points_3d = unproject_points(points_2d, points_depth, camera_matrix_tmp, normalize=normalize_points) # BxHxWx3
return points_3d.permute(0, 3, 1, 2) # Bx3xHxW
def render_gaussian2d( mean, std, size, normalized_coordinates=True):
"""
Renders the PDF of a 2D Gaussian distribution.
"""
if not (std.dtype == mean.dtype and std.device == mean.device):
raise TypeError("Expected inputs to have the same dtype and device")
height, width = size
# Create coordinates grid.
grid = create_meshgrid(height, width, normalized_coordinates, mean.device)
grid = grid.to(mean.dtype)
pos_x = grid[..., 0].view(height, width)
pos_y = grid[..., 1].view(height, width)
# dists <- (x - \mu)^2
dist_x = (pos_x - mean[..., 0, None, None]) ** 2
dist_y = (pos_y - mean[..., 1, None, None]) ** 2
# ks <- -1 / (2 \sigma^2)
k_x = -0.5 * torch.reciprocal(std[..., 0, None, None])
k_y = -0.5 * torch.reciprocal(std[..., 1, None, None])
# Assemble the 2D Gaussian.
exps_x = torch.exp(dist_x * k_x)
exps_y = torch.exp(dist_y * k_y)
gauss = exps_x * exps_y
# Rescale so that values sum to one.
val_sum = gauss.sum(-2, keepdim=True).sum(-1, keepdim=True)
gauss = _safe_zero_division(gauss, val_sum)
return gauss
def spatial_expectation2d(input, normalized_coordinates=True):
"""
Computes the expectation of coordinate values using spatial probabilities.
The input heatmap is assumed to represent a valid spatial probability
distribution, which can be achieved using `dsnt.spatial_softmax2d`.
Returns the expected value of the 2D coordinates.
The output order of the coordinates is (x, y).
"""
_validate_batched_image_tensor_input(input)
batch_size, channels, height, width = input.shape
# Create coordinates grid.
grid = create_meshgrid(height, width, normalized_coordinates, input.device)
grid = grid.to(input.dtype)
pos_x = grid[..., 0].reshape(-1)
pos_y = grid[..., 1].reshape(-1)
input_flat = input.view(batch_size, channels, -1)
# Compute the expectation of the coordinates.
expected_y = torch.sum(pos_y * input_flat, -1, keepdim=True)
expected_x = torch.sum(pos_x * input_flat, -1, keepdim=True)
output = torch.cat([expected_x, expected_y], -1)
return output.view(batch_size, channels, 2) # BxNx2
def forward(self, desc_f1, desc_f2, matches, data):
"""
Args:
feat0 (torch.Tensor): [M, WW, C]
feat1 (torch.Tensor): [M, WW, C]
data (dict)
Update:
data (dict):{
'expec_f' (torch.Tensor): [M, 3],
'mkpts0_f' (torch.Tensor): [M, 2],
'mkpts1_f' (torch.Tensor): [M, 2]}
"""
M, WW, C = desc_f1.shape
W = int(math.sqrt(WW))
scale = data['img_size_1'][0] / data['fine_size_1'][0]
self.M, self.W, self.WW, self.C, self.scale = M, W, WW, C, scale
# corner case: if no coarse matches found
if M == 0:
assert self.training == False, "M is always >0, when training, see coarse_matching.py"
# logger.warning('No matches found in coarse-level.')
matches.update({
'expec_f': torch.empty(0, 3, device=desc_f1.device),
'kpts1_f': matches['kpts1_c'],
'kpts2_f': matches['kpts2_c'],
})
return
feat_f0_picked = feat_f0_picked = desc_f1[:, WW // 2, :]
sim_matrix = torch.einsum('mc,mrc->mr', feat_f0_picked, desc_f2)
softmax_temp = 1. / C**.5
heatmap = torch.softmax(softmax_temp * sim_matrix,
dim=1).view(-1, W, W)
# compute coordinates from heatmap
coords_normalized = dsnt.spatial_expectation2d(heatmap[None],
True)[0] # [M, 2]
grid_normalized = create_meshgrid(W, W, True, heatmap.device).reshape(
1, -1, 2) # [1, WW, 2]
# compute std over <x, y>
var = torch.sum(grid_normalized**2 * heatmap.view(-1, WW, 1),
dim=1) - coords_normalized**2 # [M, 2]
std = torch.sum(torch.sqrt(torch.clamp(var, min=1e-10)),
-1) # [M] clamp needed for numerical stability
# compute absolute kpt coords
matches = self.get_fine_match(coords_normalized, matches, data)
return matches
def __init__(self,
height,
width,
mode='bilinear',
padding_mode='zeros',
normalized_coordinates=True,
align_corners=False):
super(HomographyWarper, self).__init__()
self.width = width
self.height = height
self.mode = mode
self.padding_mode = padding_mode
self.normalized_coordinates = normalized_coordinates
self.align_corners = align_corners
# create base grid to compute the flow
self.grid = create_meshgrid(
height, width, normalized_coordinates=normalized_coordinates)
# initialize the warped destination grid
self._warped_grid = None
def _create_meshgrid(height, width):
grid = create_meshgrid(height, width,
normalized_coordinates=False) # 1xHxWx2
return convert_points_to_homogeneous(grid) # append ones to last dim