File size: 3,855 Bytes
e6f80f3 | 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 | """Tensor geometry in meters, right-handed Z-up, full extents, bottom-center."""
from __future__ import annotations
import math
import torch
def wrap_yaw(angle):
return (angle + math.pi) % (2 * math.pi) - math.pi
def encode_yaw(yaw: torch.Tensor, bins: int):
if bins < 1 or not torch.isfinite(yaw).all():
raise ValueError("yaw encoding needs finite radians and positive bin count")
delta = 2 * math.pi / bins
indices = torch.floor(torch.remainder(yaw + delta / 2, 2 * math.pi) / delta).long() % bins
# Bin centres in the input dtype: long * python float would round to float32.
residual = wrap_yaw(yaw - indices.to(yaw.dtype) * delta) / (delta / 2)
tolerance = 32 * torch.finfo(yaw.dtype).eps * max(1, bins)
if ((residual < -1 - tolerance) | (residual > 1 + tolerance)).any():
raise ValueError("yaw residual outside encoding tolerance")
return indices, residual.clamp(-1, 1)
def decode_yaw(logits: torch.Tensor, residuals: torch.Tensor):
if logits.shape != residuals.shape or logits.shape[-1] < 1:
raise ValueError("yaw logits and residual shape mismatch")
k = logits.argmax(-1)
selected = residuals.gather(-1, k.unsqueeze(-1)).squeeze(-1)
# No clipping: the trained tanh head constrains residuals, recorded in model config.
# Bin centres need at least float32: a bf16 residual (autocast) would quantize them to ~0.03 rad.
dtype = torch.promote_types(selected.dtype, torch.float32)
return wrap_yaw(k.to(dtype) * (2 * math.pi / logits.shape[-1]) + selected.to(dtype) * (math.pi / logits.shape[-1]))
def encode_grid_position(xy: torch.Tensor, grid: int):
"""Row-major cell index (x_index * grid + y_index) and half-cell-unit XY residual of normalized positions.
Mirrors ``encode_yaw``: a position inside its cell has residual in [-1, 1). Positions outside
[0, 1) (bottom centres beyond the room bounding box) clamp to the border cell and a residual of +-1.
"""
if grid < 1 or xy.shape[-1] != 2 or not torch.isfinite(xy).all():
raise ValueError("grid encoding needs finite normalized XY and a positive grid size")
index = (xy * grid).floor().long().clamp(0, grid - 1)
residual = (xy - (index.to(xy.dtype) + .5) / grid) * (2 * grid)
return index[..., 0] * grid + index[..., 1], residual.clamp(-1, 1)
def decode_grid_position(logits: torch.Tensor, residuals: torch.Tensor, grid: int):
"""Argmax cell centre plus that cell's residual; at least float32 like ``decode_yaw``."""
if grid < 1 or logits.shape[-1] != grid * grid or residuals.shape != (*logits.shape, 2):
raise ValueError("grid logits must hold grid*grid cells with one XY residual pair per cell")
k = logits.argmax(-1)
selected = residuals.gather(-2, k[..., None, None].expand(*k.shape, 1, 2)).squeeze(-2)
dtype = torch.promote_types(selected.dtype, torch.float32)
index = torch.stack((torch.div(k, grid, rounding_mode="floor"), k % grid), -1).to(dtype)
return (index + .5) / grid + selected.to(dtype) / (2 * grid)
def bottom_to_center(position: torch.Tensor, size: torch.Tensor):
return position + torch.stack((torch.zeros_like(size[..., 2]), torch.zeros_like(size[..., 2]), size[..., 2] / 2), -1)
def bev_corners(position, size, yaw):
signs = size.new_tensor([[-1, -1], [1, -1], [1, 1], [-1, 1]])
local = size[..., None, :2] * signs / 2
c, s = yaw.cos()[..., None], yaw.sin()[..., None]
x = local[..., 0] * c - local[..., 1] * s
y = local[..., 0] * s + local[..., 1] * c
return torch.stack((x, y), -1) + position[..., None, :2]
def box_corners(position, size, yaw):
xy = bev_corners(position, size, yaw)
low = position[..., None, 2:3].expand(*xy.shape[:-1], 1)
high = low + size[..., None, 2:3]
return torch.cat((torch.cat((xy, low), -1), torch.cat((xy, high), -1)), -2)
|