File size: 3,682 Bytes
fd01725 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
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 # in [-1, 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.Δ # convert z value to index
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]) # y axis is up
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
|