File size: 1,372 Bytes
d547008 | 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 | import numpy as np
import torch
import torch.nn.functional as F
def th_uv_grid(H: int, W: int, **tensor_kwargs) -> torch.Tensor:
'''
:param H: int
:param W: int
:param tensor_kwargs:
:return: (H, W, 2)
'''
v, u = torch.meshgrid(torch.arange(H).to(**tensor_kwargs), torch.arange(W).to(**tensor_kwargs))
return torch.stack([u, v], dim=-1)
def depth_to_xyz(intr, depth):
'''
:param intr: shape (4,)
:param depth: shape (H, W)
:return: shape (H, W, 3)
'''
fx, fy, cx, cy = intr[0], intr[1], intr[2], intr[3]
if isinstance(depth, np.ndarray):
v, u = np.meshgrid(np.arange(depth.shape[0]), np.arange(depth.shape[1]), indexing='ij')
x = (u - cx) / fx * depth
y = (v - cy) / fy * depth
return np.stack([x, y, depth], axis=-1)
elif isinstance(depth, torch.Tensor):
tensor_kwargs = dict(device=depth.device, dtype=depth.dtype)
v, u = torch.meshgrid(torch.arange(depth.shape[0]).to(**tensor_kwargs), torch.arange(depth.shape[1]).to(**tensor_kwargs))
x = (u - cx) / fx * depth
y = (v - cy) / fy * depth
return torch.stack([x, y, depth], dim=-1)
else:
raise ValueError(f'{type(depth)=}')
def apply_SE3(SE3, pnt):
assert SE3.shape == (4, 4) and pnt.shape[-1] == 3
return (SE3[:3, :3] @ pnt[..., None])[..., 0] + SE3[:3, -1]
|