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