|
|
|
|
| import torch
|
| from torch.nn.functional import grid_sample
|
|
|
| from ..utils.geometry import from_homogeneous
|
| from .utils import make_grid
|
|
|
|
|
| class PolarProjectionDepth(torch.nn.Module):
|
| def __init__(self, z_max, ppm, scale_range, z_min=None):
|
| super().__init__()
|
| self.z_max = z_max
|
| self.Δ = Δ = 1 / ppm
|
| self.z_min = z_min = Δ if z_min is None else z_min
|
| self.scale_range = scale_range
|
| z_steps = torch.arange(z_min, z_max + Δ, Δ)
|
| self.register_buffer("depth_steps", z_steps, persistent=False)
|
|
|
| def sample_depth_scores(self, pixel_scales, camera):
|
| scale_steps = camera.f[..., None, 1] / self.depth_steps.flip(-1)
|
| log_scale_steps = torch.log2(scale_steps)
|
| scale_min, scale_max = self.scale_range
|
| log_scale_norm = (log_scale_steps - scale_min) / \
|
| (scale_max - scale_min)
|
| log_scale_norm = log_scale_norm * 2 - 1
|
|
|
| values = pixel_scales.flatten(1, 2).unsqueeze(-1)
|
| indices = log_scale_norm.unsqueeze(-1)
|
| indices = torch.stack([torch.zeros_like(indices), indices], -1)
|
| depth_scores = grid_sample(values, indices, align_corners=True)
|
| depth_scores = depth_scores.reshape(
|
| pixel_scales.shape[:-1] + (len(self.depth_steps),)
|
| )
|
| return depth_scores
|
|
|
| def forward(
|
| self,
|
| image,
|
| pixel_scales,
|
| camera,
|
| return_total_score=False,
|
| ):
|
| depth_scores = self.sample_depth_scores(pixel_scales, camera)
|
| depth_prob = torch.softmax(depth_scores, dim=1)
|
| image_polar = torch.einsum("...dhw,...hwz->...dzw", image, depth_prob)
|
| if return_total_score:
|
| cell_score = torch.logsumexp(depth_scores, dim=1, keepdim=True)
|
| return image_polar, cell_score.squeeze(1)
|
| return image_polar
|
|
|
|
|
| class CartesianProjection(torch.nn.Module):
|
| def __init__(self, z_max, x_max, ppm, z_min=None):
|
| super().__init__()
|
| self.z_max = z_max
|
| self.x_max = x_max
|
| self.Δ = Δ = 1 / ppm
|
| self.z_min = z_min = Δ if z_min is None else z_min
|
|
|
| grid_xz = make_grid(
|
| x_max * 2 + Δ, z_max, step_y=Δ, step_x=Δ, orig_y=Δ, orig_x=-x_max, y_up=True
|
| )
|
| self.register_buffer("grid_xz", grid_xz, persistent=False)
|
|
|
| def grid_to_polar(self, cam):
|
| f, c = cam.f[..., 0][..., None, None], cam.c[..., 0][..., None, None]
|
| u = from_homogeneous(self.grid_xz).squeeze(-1) * f + c
|
| z_idx = (self.grid_xz[..., 1] - self.z_min) / \
|
| self.Δ
|
| z_idx = z_idx[None].expand_as(u)
|
| grid_polar = torch.stack([u, z_idx], -1)
|
| return grid_polar
|
|
|
| def sample_from_polar(self, image_polar, valid_polar, grid_uz):
|
| size = grid_uz.new_tensor(image_polar.shape[-2:][::-1])
|
| grid_uz_norm = (grid_uz + 0.5) / size * 2 - 1
|
| grid_uz_norm = grid_uz_norm * \
|
| grid_uz.new_tensor([1, -1])
|
| image_bev = grid_sample(image_polar, grid_uz_norm, align_corners=False)
|
|
|
| if valid_polar is None:
|
| valid = torch.ones_like(image_polar[..., :1, :, :])
|
| else:
|
| valid = valid_polar.to(image_polar)[:, None]
|
| valid = grid_sample(valid, grid_uz_norm, align_corners=False)
|
| valid = valid.squeeze(1) > (1 - 1e-4)
|
|
|
| return image_bev, valid
|
|
|
| def forward(self, image_polar, valid_polar, cam):
|
| grid_uz = self.grid_to_polar(cam)
|
| image, valid = self.sample_from_polar(
|
| image_polar, valid_polar, grid_uz)
|
| return image, valid, grid_uz
|
|
|
