def reproject_disparity_to_3D(disparity_tensor: torch.Tensor, Q_matrix: torch.Tensor) -> torch.Tensor: r"""Reproject the disparity tensor to a 3D point cloud. Args: disparity_tensor: Disparity tensor of shape :math:`(B, 1, H, W)`. Q_matrix: Tensor of Q matrices of shapes :math:`(B, 4, 4)`. Returns: The 3D point cloud of shape :math:`(B, H, W, 3)` """ _check_Q_matrix(Q_matrix) _check_disparity_tensor(disparity_tensor) batch_size, rows, cols, _ = disparity_tensor.shape dtype = disparity_tensor.dtype device = disparity_tensor.device uv = create_meshgrid(rows, cols, normalized_coordinates=False, device=device, dtype=dtype) uv = uv.expand(batch_size, -1, -1, -1) v, u = torch.unbind(uv, dim=-1) v, u = torch.unsqueeze(v, -1), torch.unsqueeze(u, -1) uvd = torch.stack((u, v, disparity_tensor), 1).reshape(batch_size, 3, -1).permute(0, 2, 1) points = transform_points(Q_matrix, uvd).reshape(batch_size, rows, cols, 3) # Final check that everything went well. if not points.shape == (batch_size, rows, cols, 3): raise StereoException( f"Something went wrong in `reproject_disparity_to_3D`. Expected the final output" f"to be of shape {(batch_size, rows, cols, 3)}." f"But the computed point cloud had shape {points.shape}. " f"Please ensure input are correct. If this is an error, please submit an issue." ) return points
def render_gaussian2d( mean: torch.Tensor, std: torch.Tensor, size: Tuple[int, int], normalized_coordinates: bool = True, ): r"""Renders the PDF of a 2D Gaussian distribution. Arguments: mean (torch.Tensor): the mean location of the Gaussian to render, :math:`(\mu_x, \mu_y)`. std (torch.Tensor): the standard deviation of the Gaussian to render, :math:`(\sigma_x, \sigma_y)`. size (list): the (height, width) of the output image. normalized_coordinates: whether `mean` and `std` are assumed to use coordinates normalized in the range of [-1, 1]. Otherwise, coordinates are assumed to be in the range of the output shape. Default is True. Shape: - `mean`: :math:`(*, 2)` - `std`: :math:`(*, 2)`. Should be able to be broadcast with `mean`. - Output: :math:`(*, H, W)` """ 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: torch.Tensor = create_meshgrid(height, width, normalized_coordinates, mean.device) grid = grid.to(mean.dtype) pos_x: torch.Tensor = grid[..., 0].view(height, width) pos_y: torch.Tensor = grid[..., 1].view(height, width) # Gaussian PDF = exp(-(x - \mu)^2 / (2 \sigma^2)) # = exp(dists * ks), # where dists = (x - \mu)^2 and ks = -1 / (2 \sigma^2) # 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 forward(self, feat_f0, feat_f1, 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 = feat_f0.shape W = int(math.sqrt(WW)) scale = data['hw0_i'][0] / data['hw0_f'][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: if self.training: raise ValueError("M >0, when training, see coarse_matching.py") # logger.warning('No matches found in coarse-level.') data.update({ 'expec_f': torch.empty(0, 3, device=feat_f0.device), 'mkpts0_f': data['mkpts0_c'], 'mkpts1_f': data['mkpts1_c'], }) return feat_f0_picked = feat_f0_picked = feat_f0[:, WW // 2, :] sim_matrix = torch.einsum('mc,mrc->mr', feat_f0_picked, feat_f1) 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 # for fine-level supervision data.update( {'expec_f': torch.cat([coords_normalized, std.unsqueeze(1)], -1)}) # compute absolute kpt coords self.get_fine_match(coords_normalized, data)
def spatial_expectation2d( input: torch.Tensor, normalized_coordinates: bool = True, ) -> torch.Tensor: r"""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 :class:`~kornia.geometry.dsnt.spatial_softmax2d`. Returns the expected value of the 2D coordinates. The output order of the coordinates is (x, y). Arguments: input (torch.Tensor): the input tensor representing dense spatial probabilities. normalized_coordinates (bool): whether to return the coordinates normalized in the range of [-1, 1]. Otherwise, it will return the coordinates in the range of the input shape. Default is True. Shape: - Input: :math:`(B, N, H, W)` - Output: :math:`(B, N, 2)` Examples: >>> heatmaps = torch.tensor([[[ [0., 0., 0.], [0., 0., 0.], [0., 1., 0.]]]]) >>> coords = spatial_expectation_2d(heatmaps, False) tensor([[[1.0000, 2.0000]]]) """ _validate_batched_image_tensor_input(input) batch_size, channels, height, width = input.shape # Create coordinates grid. grid: torch.Tensor = create_meshgrid(height, width, normalized_coordinates, input.device) grid = grid.to(input.dtype) pos_x: torch.Tensor = grid[..., 0].reshape(-1) pos_y: torch.Tensor = grid[..., 1].reshape(-1) input_flat: torch.Tensor = input.view(batch_size, channels, -1) # Compute the expectation of the coordinates. expected_y: torch.Tensor = torch.sum(pos_y * input_flat, -1, keepdim=True) expected_x: torch.Tensor = torch.sum(pos_x * input_flat, -1, keepdim=True) output: torch.Tensor = torch.cat([expected_x, expected_y], -1) return output.view(batch_size, channels, 2) # BxNx2