1101 lines
36 KiB
Python
1101 lines
36 KiB
Python
import numpy as np
|
|
from typing import *
|
|
from numbers import Number
|
|
from ._helpers import batched
|
|
|
|
|
|
__all__ = [
|
|
'perspective',
|
|
'perspective_from_fov',
|
|
'perspective_from_fov_xy',
|
|
'intrinsics_from_focal_center',
|
|
'intrinsics_from_fov',
|
|
'fov_to_focal',
|
|
'focal_to_fov',
|
|
'intrinsics_to_fov',
|
|
'view_look_at',
|
|
'extrinsics_look_at',
|
|
'perspective_to_intrinsics',
|
|
'perspective_to_near_far',
|
|
'intrinsics_to_perspective',
|
|
'extrinsics_to_view',
|
|
'view_to_extrinsics',
|
|
'normalize_intrinsics',
|
|
'crop_intrinsics',
|
|
'pixel_to_uv',
|
|
'pixel_to_ndc',
|
|
'uv_to_pixel',
|
|
'project_depth',
|
|
'depth_buffer_to_linear',
|
|
'unproject_cv',
|
|
'unproject_gl',
|
|
'project_cv',
|
|
'project_gl',
|
|
'quaternion_to_matrix',
|
|
'axis_angle_to_matrix',
|
|
'matrix_to_quaternion',
|
|
'extrinsics_to_essential',
|
|
'euler_axis_angle_rotation',
|
|
'euler_angles_to_matrix',
|
|
'skew_symmetric',
|
|
'rotation_matrix_from_vectors',
|
|
'ray_intersection',
|
|
'se3_matrix',
|
|
'slerp_quaternion',
|
|
'slerp_vector',
|
|
'lerp',
|
|
'lerp_se3_matrix',
|
|
'piecewise_lerp',
|
|
'piecewise_lerp_se3_matrix',
|
|
'apply_transform'
|
|
]
|
|
|
|
|
|
@batched(0,0,0,0)
|
|
def perspective(
|
|
fov_y: Union[float, np.ndarray],
|
|
aspect: Union[float, np.ndarray],
|
|
near: Union[float, np.ndarray],
|
|
far: Union[float, np.ndarray]
|
|
) -> np.ndarray:
|
|
"""
|
|
Get OpenGL perspective matrix
|
|
|
|
Args:
|
|
fov_y (float | np.ndarray): field of view in y axis
|
|
aspect (float | np.ndarray): aspect ratio
|
|
near (float | np.ndarray): near plane to clip
|
|
far (float | np.ndarray): far plane to clip
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 4, 4] perspective matrix
|
|
"""
|
|
N = fov_y.shape[0]
|
|
ret = np.zeros((N, 4, 4), dtype=fov_y.dtype)
|
|
ret[:, 0, 0] = 1. / (np.tan(fov_y / 2) * aspect)
|
|
ret[:, 1, 1] = 1. / (np.tan(fov_y / 2))
|
|
ret[:, 2, 2] = (near + far) / (near - far)
|
|
ret[:, 2, 3] = 2. * near * far / (near - far)
|
|
ret[:, 3, 2] = -1.
|
|
return ret
|
|
|
|
|
|
def perspective_from_fov(
|
|
fov: Union[float, np.ndarray],
|
|
width: Union[int, np.ndarray],
|
|
height: Union[int, np.ndarray],
|
|
near: Union[float, np.ndarray],
|
|
far: Union[float, np.ndarray]
|
|
) -> np.ndarray:
|
|
"""
|
|
Get OpenGL perspective matrix from field of view in largest dimension
|
|
|
|
Args:
|
|
fov (float | np.ndarray): field of view in largest dimension
|
|
width (int | np.ndarray): image width
|
|
height (int | np.ndarray): image height
|
|
near (float | np.ndarray): near plane to clip
|
|
far (float | np.ndarray): far plane to clip
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 4, 4] perspective matrix
|
|
"""
|
|
fov_y = 2 * np.arctan(np.tan(fov / 2) * height / np.maximum(width, height))
|
|
aspect = width / height
|
|
return perspective(fov_y, aspect, near, far)
|
|
|
|
|
|
def perspective_from_fov_xy(
|
|
fov_x: Union[float, np.ndarray],
|
|
fov_y: Union[float, np.ndarray],
|
|
near: Union[float, np.ndarray],
|
|
far: Union[float, np.ndarray]
|
|
) -> np.ndarray:
|
|
"""
|
|
Get OpenGL perspective matrix from field of view in x and y axis
|
|
|
|
Args:
|
|
fov_x (float | np.ndarray): field of view in x axis
|
|
fov_y (float | np.ndarray): field of view in y axis
|
|
near (float | np.ndarray): near plane to clip
|
|
far (float | np.ndarray): far plane to clip
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 4, 4] perspective matrix
|
|
"""
|
|
aspect = np.tan(fov_x / 2) / np.tan(fov_y / 2)
|
|
return perspective(fov_y, aspect, near, far)
|
|
|
|
|
|
def intrinsics_from_focal_center(
|
|
fx: Union[float, np.ndarray],
|
|
fy: Union[float, np.ndarray],
|
|
cx: Union[float, np.ndarray],
|
|
cy: Union[float, np.ndarray],
|
|
dtype: Optional[np.dtype] = np.float32
|
|
) -> np.ndarray:
|
|
"""
|
|
Get OpenCV intrinsics matrix
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 3, 3] OpenCV intrinsics matrix
|
|
"""
|
|
if any(isinstance(x, np.ndarray) for x in (fx, fy, cx, cy)):
|
|
dtype = np.result_type(fx, fy, cx, cy)
|
|
fx, fy, cx, cy = np.broadcast_arrays(fx, fy, cx, cy)
|
|
ret = np.zeros((*fx.shape, 3, 3), dtype=dtype)
|
|
ret[..., 0, 0] = fx
|
|
ret[..., 1, 1] = fy
|
|
ret[..., 0, 2] = cx
|
|
ret[..., 1, 2] = cy
|
|
ret[..., 2, 2] = 1.
|
|
return ret
|
|
|
|
|
|
def intrinsics_from_fov(
|
|
fov_max: Union[float, np.ndarray] = None,
|
|
fov_min: Union[float, np.ndarray] = None,
|
|
fov_x: Union[float, np.ndarray] = None,
|
|
fov_y: Union[float, np.ndarray] = None,
|
|
width: Union[int, np.ndarray] = None,
|
|
height: Union[int, np.ndarray] = None,
|
|
) -> np.ndarray:
|
|
"""
|
|
Get normalized OpenCV intrinsics matrix from given field of view.
|
|
You can provide either fov_max, fov_min, fov_x or fov_y
|
|
|
|
Args:
|
|
width (int | np.ndarray): image width
|
|
height (int | np.ndarray): image height
|
|
fov_max (float | np.ndarray): field of view in largest dimension
|
|
fov_min (float | np.ndarray): field of view in smallest dimension
|
|
fov_x (float | np.ndarray): field of view in x axis
|
|
fov_y (float | np.ndarray): field of view in y axis
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 3, 3] OpenCV intrinsics matrix
|
|
"""
|
|
if fov_max is not None:
|
|
fx = np.maximum(width, height) / width / (2 * np.tan(fov_max / 2))
|
|
fy = np.maximum(width, height) / height / (2 * np.tan(fov_max / 2))
|
|
elif fov_min is not None:
|
|
fx = np.minimum(width, height) / width / (2 * np.tan(fov_min / 2))
|
|
fy = np.minimum(width, height) / height / (2 * np.tan(fov_min / 2))
|
|
elif fov_x is not None and fov_y is not None:
|
|
fx = 1 / (2 * np.tan(fov_x / 2))
|
|
fy = 1 / (2 * np.tan(fov_y / 2))
|
|
elif fov_x is not None:
|
|
fx = 1 / (2 * np.tan(fov_x / 2))
|
|
fy = fx * width / height
|
|
elif fov_y is not None:
|
|
fy = 1 / (2 * np.tan(fov_y / 2))
|
|
fx = fy * height / width
|
|
cx = 0.5
|
|
cy = 0.5
|
|
ret = intrinsics_from_focal_center(fx, fy, cx, cy)
|
|
return ret
|
|
|
|
|
|
def focal_to_fov(focal: np.ndarray):
|
|
return 2 * np.arctan(0.5 / focal)
|
|
|
|
|
|
def fov_to_focal(fov: np.ndarray):
|
|
return 0.5 / np.tan(fov / 2)
|
|
|
|
|
|
def intrinsics_to_fov(intrinsics: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
|
fov_x = focal_to_fov(intrinsics[..., 0, 0])
|
|
fov_y = focal_to_fov(intrinsics[..., 1, 1])
|
|
return fov_x, fov_y
|
|
|
|
|
|
@batched(1,1,1)
|
|
def view_look_at(
|
|
eye: np.ndarray,
|
|
look_at: np.ndarray,
|
|
up: np.ndarray
|
|
) -> np.ndarray:
|
|
"""
|
|
Get OpenGL view matrix looking at something
|
|
|
|
Args:
|
|
eye (np.ndarray): [..., 3] the eye position
|
|
look_at (np.ndarray): [..., 3] the position to look at
|
|
up (np.ndarray): [..., 3] head up direction (y axis in screen space). Not necessarily othogonal to view direction
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 4, 4], view matrix
|
|
"""
|
|
z = eye - look_at
|
|
x = np.cross(up, z)
|
|
y = np.cross(z, x)
|
|
# x = np.cross(y, z)
|
|
x = x / np.linalg.norm(x, axis=-1, keepdims=True)
|
|
y = y / np.linalg.norm(y, axis=-1, keepdims=True)
|
|
z = z / np.linalg.norm(z, axis=-1, keepdims=True)
|
|
R = np.stack([x, y, z], axis=-2)
|
|
t = -np.matmul(R, eye[..., None])
|
|
return np.concatenate([
|
|
np.concatenate([R, t], axis=-1),
|
|
np.array([[[0., 0., 0., 1.]]]).repeat(eye.shape[0], axis=0)
|
|
], axis=-2)
|
|
|
|
|
|
@batched(1,1,1)
|
|
def extrinsics_look_at(
|
|
eye: np.ndarray,
|
|
look_at: np.ndarray,
|
|
up: np.ndarray
|
|
) -> np.ndarray:
|
|
"""
|
|
Get OpenCV extrinsics matrix looking at something
|
|
|
|
Args:
|
|
eye (np.ndarray): [..., 3] the eye position
|
|
look_at (np.ndarray): [..., 3] the position to look at
|
|
up (np.ndarray): [..., 3] head up direction (-y axis in screen space). Not necessarily othogonal to view direction
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 4, 4], extrinsics matrix
|
|
"""
|
|
z = look_at - eye
|
|
x = np.cross(-up, z)
|
|
y = np.cross(z, x)
|
|
# x = np.cross(y, z)
|
|
x = x / np.linalg.norm(x, axis=-1, keepdims=True)
|
|
y = y / np.linalg.norm(y, axis=-1, keepdims=True)
|
|
z = z / np.linalg.norm(z, axis=-1, keepdims=True)
|
|
R = np.stack([x, y, z], axis=-2)
|
|
t = -np.matmul(R, eye[..., None])
|
|
return np.concatenate([
|
|
np.concatenate([R, t], axis=-1),
|
|
np.array([[[0., 0., 0., 1.]]], dtype=eye.dtype).repeat(eye.shape[0], axis=0)
|
|
], axis=-2)
|
|
|
|
|
|
def perspective_to_intrinsics(
|
|
perspective: np.ndarray
|
|
) -> np.ndarray:
|
|
"""
|
|
OpenGL perspective matrix to OpenCV intrinsics
|
|
|
|
Args:
|
|
perspective (np.ndarray): [..., 4, 4] OpenGL perspective matrix
|
|
|
|
Returns:
|
|
(np.ndarray): shape [..., 3, 3] OpenCV intrinsics
|
|
"""
|
|
ret = np.array([[0.5, 0., 0.5], [0., -0.5, 0.5], [0., 0., 1.]], dtype=perspective.dtype) \
|
|
@ perspective[..., [0, 1, 3], :3] \
|
|
@ np.diag(np.array([1, -1, -1], dtype=perspective.dtype))
|
|
return ret
|
|
|
|
|
|
def perspective_to_near_far(perspective: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""
|
|
Get near and far planes from OpenGL perspective matrix
|
|
|
|
Args:
|
|
"""
|
|
a, b = perspective[..., 2, 2], perspective[..., 2, 3]
|
|
near, far = b / (a - 1), b / (a + 1)
|
|
return near, far
|
|
|
|
|
|
@batched(2,0,0)
|
|
def intrinsics_to_perspective(
|
|
intrinsics: np.ndarray,
|
|
near: Union[float, np.ndarray],
|
|
far: Union[float, np.ndarray],
|
|
) -> np.ndarray:
|
|
"""
|
|
OpenCV intrinsics to OpenGL perspective matrix
|
|
NOTE: not work for tile-shifting intrinsics currently
|
|
|
|
Args:
|
|
intrinsics (np.ndarray): [..., 3, 3] OpenCV intrinsics matrix
|
|
near (float | np.ndarray): [...] near plane to clip
|
|
far (float | np.ndarray): [...] far plane to clip
|
|
Returns:
|
|
(np.ndarray): [..., 4, 4] OpenGL perspective matrix
|
|
"""
|
|
N = intrinsics.shape[0]
|
|
fx, fy = intrinsics[:, 0, 0], intrinsics[:, 1, 1]
|
|
cx, cy = intrinsics[:, 0, 2], intrinsics[:, 1, 2]
|
|
ret = np.zeros((N, 4, 4), dtype=intrinsics.dtype)
|
|
ret[:, 0, 0] = 2 * fx
|
|
ret[:, 1, 1] = 2 * fy
|
|
ret[:, 0, 2] = -2 * cx + 1
|
|
ret[:, 1, 2] = 2 * cy - 1
|
|
ret[:, 2, 2] = (near + far) / (near - far)
|
|
ret[:, 2, 3] = 2. * near * far / (near - far)
|
|
ret[:, 3, 2] = -1.
|
|
return ret
|
|
|
|
|
|
@batched(2)
|
|
def extrinsics_to_view(
|
|
extrinsics: np.ndarray
|
|
) -> np.ndarray:
|
|
"""
|
|
OpenCV camera extrinsics to OpenGL view matrix
|
|
|
|
Args:
|
|
extrinsics (np.ndarray): [..., 4, 4] OpenCV camera extrinsics matrix
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 4, 4] OpenGL view matrix
|
|
"""
|
|
return extrinsics * np.array([1, -1, -1, 1], dtype=extrinsics.dtype)[:, None]
|
|
|
|
|
|
@batched(2)
|
|
def view_to_extrinsics(
|
|
view: np.ndarray
|
|
) -> np.ndarray:
|
|
"""
|
|
OpenGL view matrix to OpenCV camera extrinsics
|
|
|
|
Args:
|
|
view (np.ndarray): [..., 4, 4] OpenGL view matrix
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 4, 4] OpenCV camera extrinsics matrix
|
|
"""
|
|
return view * np.array([1, -1, -1, 1], dtype=view.dtype)[:, None]
|
|
|
|
|
|
@batched(2, 0, 0, None)
|
|
def normalize_intrinsics(
|
|
intrinsics: np.ndarray,
|
|
width: Union[int, np.ndarray],
|
|
height: Union[int, np.ndarray],
|
|
integer_pixel_centers: bool = True
|
|
) -> np.ndarray:
|
|
"""
|
|
Normalize intrinsics from pixel cooridnates to uv coordinates
|
|
|
|
Args:
|
|
intrinsics (np.ndarray): [..., 3, 3] camera intrinsics(s) to normalize
|
|
width (int | np.ndarray): [...] image width(s)
|
|
height (int | np.ndarray): [...] image height(s)
|
|
integer_pixel_centers (bool): whether the integer pixel coordinates are at the center of the pixel. If False, the integer coordinates are at the left-top corner of the pixel.
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 3, 3] normalized camera intrinsics(s)
|
|
"""
|
|
zeros = np.zeros_like(width)
|
|
ones = np.ones_like(width)
|
|
if integer_pixel_centers:
|
|
transform = np.stack([
|
|
1 / width, zeros, 0.5 / width,
|
|
zeros, 1 / height, 0.5 / height,
|
|
zeros, zeros, ones
|
|
]).reshape(*zeros.shape, 3, 3)
|
|
else:
|
|
transform = np.stack([
|
|
1 / width, zeros, zeros,
|
|
zeros, 1 / height, zeros,
|
|
zeros, zeros, ones
|
|
]).reshape(*zeros.shape, 3, 3)
|
|
return transform @ intrinsics
|
|
|
|
|
|
@batched(2,0,0,0,0,0,0)
|
|
def crop_intrinsics(
|
|
intrinsics: np.ndarray,
|
|
width: Union[int, np.ndarray],
|
|
height: Union[int, np.ndarray],
|
|
left: Union[int, np.ndarray],
|
|
top: Union[int, np.ndarray],
|
|
crop_width: Union[int, np.ndarray],
|
|
crop_height: Union[int, np.ndarray]
|
|
) -> np.ndarray:
|
|
"""
|
|
Evaluate the new intrinsics(s) after crop the image: cropped_img = img[top:top+crop_height, left:left+crop_width]
|
|
|
|
Args:
|
|
intrinsics (np.ndarray): [..., 3, 3] camera intrinsics(s) to crop
|
|
width (int | np.ndarray): [...] image width(s)
|
|
height (int | np.ndarray): [...] image height(s)
|
|
left (int | np.ndarray): [...] left crop boundary
|
|
top (int | np.ndarray): [...] top crop boundary
|
|
crop_width (int | np.ndarray): [...] crop width
|
|
crop_height (int | np.ndarray): [...] crop height
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 3, 3] cropped camera intrinsics(s)
|
|
"""
|
|
zeros = np.zeros_like(width)
|
|
ones = np.ones_like(width)
|
|
transform = np.stack([
|
|
width / crop_width, zeros, -left / crop_width,
|
|
zeros, height / crop_height, -top / crop_height,
|
|
zeros, zeros, ones
|
|
]).reshape(*zeros.shape, 3, 3)
|
|
return transform @ intrinsics
|
|
|
|
|
|
@batched(1,0,0)
|
|
def pixel_to_uv(
|
|
pixel: np.ndarray,
|
|
width: Union[int, np.ndarray],
|
|
height: Union[int, np.ndarray]
|
|
) -> np.ndarray:
|
|
"""
|
|
Args:
|
|
pixel (np.ndarray): [..., 2] pixel coordinrates defined in image space, x range is (0, W - 1), y range is (0, H - 1)
|
|
width (int | np.ndarray): [...] image width(s)
|
|
height (int | np.ndarray): [...] image height(s)
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 2] pixel coordinrates defined in uv space, the range is (0, 1)
|
|
"""
|
|
if not np.issubdtype(pixel.dtype, np.floating):
|
|
pixel = pixel.astype(np.float32)
|
|
dtype = pixel.dtype
|
|
uv = (pixel + np.array(0.5, dtype=dtype)) / np.stack([width, height], axis=-1)
|
|
return uv
|
|
|
|
|
|
@batched(1,0,0)
|
|
def uv_to_pixel(
|
|
uv: np.ndarray,
|
|
width: Union[int, np.ndarray],
|
|
height: Union[int, np.ndarray]
|
|
) -> np.ndarray:
|
|
"""
|
|
Args:
|
|
pixel (np.ndarray): [..., 2] pixel coordinrates defined in image space, x range is (0, W - 1), y range is (0, H - 1)
|
|
width (int | np.ndarray): [...] image width(s)
|
|
height (int | np.ndarray): [...] image height(s)
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 2] pixel coordinrates defined in uv space, the range is (0, 1)
|
|
"""
|
|
pixel = uv * np.stack([width, height], axis=-1) - 0.5
|
|
return pixel
|
|
|
|
|
|
@batched(1,0,0)
|
|
def pixel_to_ndc(
|
|
pixel: np.ndarray,
|
|
width: Union[int, np.ndarray],
|
|
height: Union[int, np.ndarray]
|
|
) -> np.ndarray:
|
|
"""
|
|
Args:
|
|
pixel (np.ndarray): [..., 2] pixel coordinrates defined in image space, x range is (0, W - 1), y range is (0, H - 1)
|
|
width (int | np.ndarray): [...] image width(s)
|
|
height (int | np.ndarray): [...] image height(s)
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 2] pixel coordinrates defined in ndc space, the range is (-1, 1)
|
|
"""
|
|
if not np.issubdtype(pixel.dtype, np.floating):
|
|
pixel = pixel.astype(np.float32)
|
|
dtype = pixel.dtype
|
|
ndc = (pixel + np.array(0.5, dtype=dtype)) / (np.stack([width, height], dim=-1) * np.array([2, -2], dtype=dtype)) \
|
|
+ np.array([-1, 1], dtype=dtype)
|
|
return ndc
|
|
|
|
|
|
@batched(0,0,0)
|
|
def project_depth(
|
|
depth: np.ndarray,
|
|
near: Union[float, np.ndarray],
|
|
far: Union[float, np.ndarray]
|
|
) -> np.ndarray:
|
|
"""
|
|
Project linear depth to depth value in screen space
|
|
|
|
Args:
|
|
depth (np.ndarray): [...] depth value
|
|
near (float | np.ndarray): [...] near plane to clip
|
|
far (float | np.ndarray): [...] far plane to clip
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 1] depth value in screen space, value ranging in [0, 1]
|
|
"""
|
|
return (far - near * far / depth) / (far - near)
|
|
|
|
|
|
@batched(0,0,0)
|
|
def depth_buffer_to_linear(
|
|
depth_buffer: np.ndarray,
|
|
near: Union[float, np.ndarray],
|
|
far: Union[float, np.ndarray]
|
|
) -> np.ndarray:
|
|
"""
|
|
OpenGL depth buffer to linear depth
|
|
|
|
Args:
|
|
depth_buffer (np.ndarray): [...] depth value
|
|
near (float | np.ndarray): [...] near plane to clip
|
|
far (float | np.ndarray): [...] far plane to clip
|
|
|
|
Returns:
|
|
(np.ndarray): [..., 1] linear depth
|
|
"""
|
|
return near * far / (far - (far - near) * depth_buffer)
|
|
|
|
|
|
@batched(2,2,2,2)
|
|
def project_gl(
|
|
points: np.ndarray,
|
|
model: np.ndarray = None,
|
|
view: np.ndarray = None,
|
|
perspective: np.ndarray = None
|
|
) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""
|
|
Project 3D points to 2D following the OpenGL convention (except for row major matrice)
|
|
|
|
Args:
|
|
points (np.ndarray): [..., N, 3] or [..., N, 4] 3D points to project, if the last
|
|
dimension is 4, the points are assumed to be in homogeneous coordinates
|
|
model (np.ndarray): [..., 4, 4] model matrix
|
|
view (np.ndarray): [..., 4, 4] view matrix
|
|
perspective (np.ndarray): [..., 4, 4] perspective matrix
|
|
|
|
Returns:
|
|
scr_coord (np.ndarray): [..., N, 3] screen space coordinates, value ranging in [0, 1].
|
|
The origin (0., 0., 0.) is corresponding to the left & bottom & nearest
|
|
linear_depth (np.ndarray): [..., N] linear depth
|
|
"""
|
|
assert perspective is not None, "perspective matrix is required"
|
|
if points.shape[-1] == 3:
|
|
points = np.concatenate([points, np.ones_like(points[..., :1])], axis=-1)
|
|
if model is not None:
|
|
points = points @ model.swapaxes(-1, -2)
|
|
if view is not None:
|
|
points = points @ view.swapaxes(-1, -2)
|
|
clip_coord = points @ perspective.swapaxes(-1, -2)
|
|
ndc_coord = clip_coord[..., :3] / clip_coord[..., 3:]
|
|
scr_coord = ndc_coord * 0.5 + 0.5
|
|
linear_depth = clip_coord[..., 3]
|
|
return scr_coord, linear_depth
|
|
|
|
|
|
@batched(2,2,2)
|
|
def project_cv(
|
|
points: np.ndarray,
|
|
extrinsics: np.ndarray = None,
|
|
intrinsics: np.ndarray = None
|
|
) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""
|
|
Project 3D points to 2D following the OpenCV convention
|
|
|
|
Args:
|
|
points (np.ndarray): [..., N, 3] or [..., N, 4] 3D points to project, if the last
|
|
dimension is 4, the points are assumed to be in homogeneous coordinates
|
|
extrinsics (np.ndarray): [..., 4, 4] extrinsics matrix
|
|
intrinsics (np.ndarray): [..., 3, 3] intrinsics matrix
|
|
|
|
Returns:
|
|
uv_coord (np.ndarray): [..., N, 2] uv coordinates, value ranging in [0, 1].
|
|
The origin (0., 0.) is corresponding to the left & top
|
|
linear_depth (np.ndarray): [..., N] linear depth
|
|
"""
|
|
assert intrinsics is not None, "intrinsics matrix is required"
|
|
if points.shape[-1] == 3:
|
|
points = np.concatenate([points, np.ones_like(points[..., :1])], axis=-1)
|
|
if extrinsics is not None:
|
|
points = points @ extrinsics.swapaxes(-1, -2)
|
|
points = points[..., :3] @ intrinsics.swapaxes(-1, -2)
|
|
uv_coord = points[..., :2] / points[..., 2:]
|
|
linear_depth = points[..., 2]
|
|
return uv_coord, linear_depth
|
|
|
|
|
|
@batched(2,2,2,2)
|
|
def unproject_gl(
|
|
screen_coord: np.ndarray,
|
|
model: np.ndarray = None,
|
|
view: np.ndarray = None,
|
|
perspective: np.ndarray = None
|
|
) -> np.ndarray:
|
|
"""
|
|
Unproject screen space coordinates to 3D view space following the OpenGL convention (except for row major matrice)
|
|
|
|
Args:
|
|
screen_coord (np.ndarray): [..., N, 3] screen space coordinates, value ranging in [0, 1].
|
|
The origin (0., 0., 0.) is corresponding to the left & bottom & nearest
|
|
model (np.ndarray): [..., 4, 4] model matrix
|
|
view (np.ndarray): [..., 4, 4] view matrix
|
|
perspective (np.ndarray): [..., 4, 4] perspective matrix
|
|
|
|
Returns:
|
|
points (np.ndarray): [..., N, 3] 3d points
|
|
"""
|
|
assert perspective is not None, "perspective matrix is required"
|
|
ndc_xy = screen_coord * 2 - 1
|
|
clip_coord = np.concatenate([ndc_xy, np.ones_like(ndc_xy[..., :1])], axis=-1)
|
|
transform = perspective
|
|
if view is not None:
|
|
transform = transform @ view
|
|
if model is not None:
|
|
transform = transform @ model
|
|
transform = np.linalg.inv(transform)
|
|
points = clip_coord @ transform.swapaxes(-1, -2)
|
|
points = points[..., :3] / points[..., 3:]
|
|
return points
|
|
|
|
|
|
@batched(2,1,2,2)
|
|
def unproject_cv(
|
|
uv_coord: np.ndarray,
|
|
depth: np.ndarray,
|
|
extrinsics: np.ndarray = None,
|
|
intrinsics: np.ndarray = None
|
|
) -> np.ndarray:
|
|
"""
|
|
Unproject uv coordinates to 3D view space following the OpenCV convention
|
|
|
|
Args:
|
|
uv_coord (np.ndarray): [..., N, 2] uv coordinates, value ranging in [0, 1].
|
|
The origin (0., 0.) is corresponding to the left & top
|
|
depth (np.ndarray): [..., N] depth value
|
|
extrinsics (np.ndarray): [..., 4, 4] extrinsics matrix
|
|
intrinsics (np.ndarray): [..., 3, 3] intrinsics matrix
|
|
|
|
Returns:
|
|
points (np.ndarray): [..., N, 3] 3d points
|
|
"""
|
|
assert intrinsics is not None, "intrinsics matrix is required"
|
|
points = np.concatenate([uv_coord, np.ones_like(uv_coord[..., :1])], axis=-1)
|
|
points = points @ np.linalg.inv(intrinsics).swapaxes(-1, -2)
|
|
points = points * depth[..., None]
|
|
if extrinsics is not None:
|
|
points = np.concatenate([points, np.ones_like(points[..., :1])], axis=-1)
|
|
points = (points @ np.linalg.inv(extrinsics).swapaxes(-1, -2))[..., :3]
|
|
return points
|
|
|
|
|
|
def quaternion_to_matrix(quaternion: np.ndarray, eps: float = 1e-12) -> np.ndarray:
|
|
"""Converts a batch of quaternions (w, x, y, z) to rotation matrices
|
|
|
|
Args:
|
|
quaternion (np.ndarray): shape (..., 4), the quaternions to convert
|
|
|
|
Returns:
|
|
np.ndarray: shape (..., 3, 3), the rotation matrices corresponding to the given quaternions
|
|
"""
|
|
assert quaternion.shape[-1] == 4
|
|
quaternion = quaternion / np.linalg.norm(quaternion, axis=-1, keepdims=True).clip(min=eps)
|
|
w, x, y, z = quaternion[..., 0], quaternion[..., 1], quaternion[..., 2], quaternion[..., 3]
|
|
zeros = np.zeros_like(w)
|
|
I = np.eye(3, dtype=quaternion.dtype)
|
|
xyz = quaternion[..., 1:]
|
|
A = xyz[..., :, None] * xyz[..., None, :] - I * (xyz ** 2).sum(axis=-1)[..., None, None]
|
|
B = np.stack([
|
|
zeros, -z, y,
|
|
z, zeros, -x,
|
|
-y, x, zeros
|
|
], axis=-1).reshape(*quaternion.shape[:-1], 3, 3)
|
|
rot_mat = I + 2 * (A + w[..., None, None] * B)
|
|
return rot_mat
|
|
|
|
|
|
def matrix_to_quaternion(rot_mat: np.ndarray, eps: float = 1e-12) -> np.ndarray:
|
|
"""Convert 3x3 rotation matrix to quaternion (w, x, y, z)
|
|
|
|
Args:
|
|
rot_mat (np.ndarray): shape (..., 3, 3), the rotation matrices to convert
|
|
|
|
Returns:
|
|
np.ndarray: shape (..., 4), the quaternions corresponding to the given rotation matrices
|
|
"""
|
|
# Extract the diagonal and off-diagonal elements of the rotation matrix
|
|
m00, m01, m02, m10, m11, m12, m20, m21, m22 = [rot_mat[..., i, j] for i in range(3) for j in range(3)]
|
|
|
|
diag = np.diagonal(rot_mat, axis1=-2, axis2=-1)
|
|
M = np.array([
|
|
[1, 1, 1],
|
|
[1, -1, -1],
|
|
[-1, 1, -1],
|
|
[-1, -1, 1]
|
|
], dtype=rot_mat.dtype)
|
|
wxyz = 0.5 * np.clip(1 + diag @ M.T, 0.0, None) ** 0.5
|
|
max_idx = np.argmax(wxyz, axis=-1)
|
|
xw = np.sign(m21 - m12)
|
|
yw = np.sign(m02 - m20)
|
|
zw = np.sign(m10 - m01)
|
|
yz = np.sign(m21 + m12)
|
|
xz = np.sign(m02 + m20)
|
|
xy = np.sign(m01 + m10)
|
|
ones = np.ones_like(xw)
|
|
sign = np.where(
|
|
max_idx[..., None] == 0,
|
|
np.stack([ones, xw, yw, zw], axis=-1),
|
|
np.where(
|
|
max_idx[..., None] == 1,
|
|
np.stack([xw, ones, xy, xz], axis=-1),
|
|
np.where(
|
|
max_idx[..., None] == 2,
|
|
np.stack([yw, xy, ones, yz], axis=-1),
|
|
np.stack([zw, xz, yz, ones], axis=-1)
|
|
)
|
|
)
|
|
)
|
|
quat = sign * wxyz
|
|
quat = quat / np.linalg.norm(quat, axis=-1, keepdims=True).clip(min=eps)
|
|
return quat
|
|
|
|
|
|
def extrinsics_to_essential(extrinsics: np.ndarray):
|
|
"""
|
|
extrinsics matrix `[[R, t] [0, 0, 0, 1]]` such that `x' = R (x - t)` to essential matrix such that `x' E x = 0`
|
|
|
|
Args:
|
|
extrinsics (np.ndaray): [..., 4, 4] extrinsics matrix
|
|
|
|
Returns:
|
|
(np.ndaray): [..., 3, 3] essential matrix
|
|
"""
|
|
assert extrinsics.shape[-2:] == (4, 4)
|
|
R = extrinsics[..., :3, :3]
|
|
t = extrinsics[..., :3, 3]
|
|
zeros = np.zeros_like(t[..., 0])
|
|
t_x = np.stack([
|
|
zeros, -t[..., 2], t[..., 1],
|
|
t[..., 2], zeros, -t[..., 0],
|
|
-t[..., 1], t[..., 0], zeros
|
|
]).reshape(*t.shape[:-1], 3, 3)
|
|
return t_x @ R
|
|
|
|
|
|
def euler_axis_angle_rotation(axis: str, angle: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Return the rotation matrices for one of the rotations about an axis
|
|
of which Euler angles describe, for each value of the angle given.
|
|
|
|
Args:
|
|
axis: Axis label "X" or "Y or "Z".
|
|
angle: any shape tensor of Euler angles in radians
|
|
|
|
Returns:
|
|
Rotation matrices as tensor of shape (..., 3, 3).
|
|
"""
|
|
|
|
cos = np.cos(angle)
|
|
sin = np.sin(angle)
|
|
one = np.ones_like(angle)
|
|
zero = np.zeros_like(angle)
|
|
|
|
if axis == "X":
|
|
R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
|
|
elif axis == "Y":
|
|
R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
|
|
elif axis == "Z":
|
|
R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
|
|
else:
|
|
raise ValueError("letter must be either X, Y or Z.")
|
|
|
|
return np.stack(R_flat, -1).reshape(angle.shape + (3, 3))
|
|
|
|
|
|
def euler_angles_to_matrix(euler_angles: np.ndarray, convention: str = 'XYZ') -> np.ndarray:
|
|
"""
|
|
Convert rotations given as Euler angles in radians to rotation matrices.
|
|
|
|
Args:
|
|
euler_angles: Euler angles in radians as ndarray of shape (..., 3), XYZ
|
|
convention: permutation of "X", "Y" or "Z", representing the order of Euler rotations to apply.
|
|
|
|
Returns:
|
|
Rotation matrices as ndarray of shape (..., 3, 3).
|
|
"""
|
|
if euler_angles.shape[-1] != 3:
|
|
raise ValueError("Invalid input euler angles.")
|
|
if len(convention) != 3:
|
|
raise ValueError("Convention must have 3 letters.")
|
|
if convention[1] in (convention[0], convention[2]):
|
|
raise ValueError(f"Invalid convention {convention}.")
|
|
for letter in convention:
|
|
if letter not in ("X", "Y", "Z"):
|
|
raise ValueError(f"Invalid letter {letter} in convention string.")
|
|
matrices = [
|
|
euler_axis_angle_rotation(c, euler_angles[..., 'XYZ'.index(c)])
|
|
for c in convention
|
|
]
|
|
return matrices[2] @ matrices[1] @ matrices[0]
|
|
|
|
|
|
def skew_symmetric(v: np.ndarray):
|
|
"Skew symmetric matrix from a 3D vector"
|
|
assert v.shape[-1] == 3, "v must be 3D"
|
|
x, y, z = v[..., 0], v[..., 1], v[..., 2]
|
|
zeros = np.zeros_like(x)
|
|
return np.stack([
|
|
zeros, -z, y,
|
|
z, zeros, -x,
|
|
-y, x, zeros,
|
|
], axis=-1).reshape(*v.shape[:-1], 3, 3)
|
|
|
|
|
|
def rotation_matrix_from_vectors(v1: np.ndarray, v2: np.ndarray):
|
|
"Rotation matrix that rotates v1 to v2"
|
|
I = np.eye(3, dtype=v1.dtype)
|
|
v1 = v1 / np.linalg.norm(v1, axis=-1)
|
|
v2 = v2 / np.linalg.norm(v2, axis=-1)
|
|
v = np.cross(v1, v2, axis=-1)
|
|
c = np.sum(v1 * v2, axis=-1)
|
|
K = skew_symmetric(v)
|
|
R = I + K + (1 / (1 + c)).astype(v1.dtype)[None, None] * (K @ K) # Avoid numpy's default type casting for scalars
|
|
return R
|
|
|
|
|
|
def axis_angle_to_matrix(axis_angle: np.ndarray, eps: float = 1e-12) -> np.ndarray:
|
|
"""Convert axis-angle representation (rotation vector) to rotation matrix, whose direction is the axis of rotation and length is the angle of rotation
|
|
|
|
Args:
|
|
axis_angle (np.ndarray): shape (..., 3), axis-angle vcetors
|
|
|
|
Returns:
|
|
np.ndarray: shape (..., 3, 3) The rotation matrices for the given axis-angle parameters
|
|
"""
|
|
batch_shape = axis_angle.shape[:-1]
|
|
dtype = axis_angle.dtype
|
|
|
|
angle = np.linalg.norm(axis_angle, axis=-1, keepdims=True)
|
|
axis = axis_angle / (angle + eps)
|
|
|
|
cos = np.cos(angle)[..., None, :]
|
|
sin = np.sin(angle)[..., None, :]
|
|
|
|
rx, ry, rz = np.split(axis, 3, axis=-1)
|
|
zeros = np.zeros((*batch_shape, 1), dtype=dtype)
|
|
K = np.concatenate([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], axis=-1).reshape((*batch_shape, 3, 3))
|
|
|
|
ident = np.eye(3, dtype=dtype)
|
|
rot_mat = ident + sin * K + (1 - cos) * (K @ K)
|
|
return rot_mat
|
|
|
|
|
|
def ray_intersection(p1: np.ndarray, d1: np.ndarray, p2: np.ndarray, d2: np.ndarray):
|
|
"""
|
|
Compute the intersection/closest point of two D-dimensional rays
|
|
If the rays are intersecting, the closest point is the intersection point.
|
|
|
|
Args:
|
|
p1 (np.ndarray): (..., D) origin of ray 1
|
|
d1 (np.ndarray): (..., D) direction of ray 1
|
|
p2 (np.ndarray): (..., D) origin of ray 2
|
|
d2 (np.ndarray): (..., D) direction of ray 2
|
|
|
|
Returns:
|
|
(np.ndarray): (..., N) intersection point
|
|
"""
|
|
p1, d1, p2, d2 = np.broadcast_arrays(p1, d1, p2, d2)
|
|
dtype = p1.dtype
|
|
dim = p1.shape[-1]
|
|
d = np.stack([d1, d2], axis=-2) # (..., 2, D)
|
|
p = np.stack([p1, p2], axis=-2) # (..., 2, D)
|
|
A = np.concatenate([
|
|
(np.eye(dim, dtype=dtype) * np.ones((*p.shape[:-2], 2, 1, 1))).reshape(*d.shape[:-2], 2 * dim, dim), # (..., 2 * D, D)
|
|
-(np.eye(2, dtype=dtype)[..., None] * d[..., None, :]).swapaxes(-2, -1).reshape(*d.shape[:-2], 2 * dim, 2) # (..., 2 * D, 2)
|
|
], axis=-1) # (..., 2 * D, D + 2)
|
|
b = p.reshape(*p.shape[:-2], 2 * dim) # (..., 2 * D)
|
|
x = np.linalg.solve(A.swapaxes(-1, -2) @ A + 1e-12 * np.eye(dim + 2, dtype=dtype), (A.swapaxes(-1, -2) @ b[..., :, None])[..., 0])
|
|
return x[..., :dim], (x[..., dim], x[..., dim + 1])
|
|
|
|
|
|
def se3_matrix(R: np.ndarray, t: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Convert rotation matrix and translation vector to 4x4 transformation matrix.
|
|
|
|
Args:
|
|
R (np.ndarray): [..., 3, 3] rotation matrix
|
|
t (np.ndarray): [..., 3] translation vector
|
|
|
|
Returns:
|
|
np.ndarray: [..., 4, 4] transformation matrix
|
|
"""
|
|
assert R.shape[:-2] == t.shape[:-1]
|
|
assert R.shape[-1] == 3 and R.shape[-2] == 3
|
|
return np.concatenate([
|
|
np.concatenate([R, t[..., None]], axis=-1),
|
|
np.concatenate([np.zeros_like(t), np.ones_like(t[..., :1])], axis=-1)[..., None, :]
|
|
], axis=-2)
|
|
|
|
|
|
def slerp_quaternion(q1: np.ndarray, q2: np.ndarray, t: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Spherical linear interpolation between two unit quaternions.
|
|
|
|
Args:
|
|
q1 (np.ndarray): [..., d] unit vector 1
|
|
q2 (np.ndarray): [..., d] unit vector 2
|
|
t (np.ndarray): [...] interpolation parameter in [0, 1]
|
|
|
|
Returns:
|
|
np.ndarray: [..., 3] interpolated unit vector
|
|
"""
|
|
q1 = q1 / np.linalg.norm(q1, axis=-1, keepdims=True)
|
|
q2 = q2 / np.linalg.norm(q2, axis=-1, keepdims=True)
|
|
dot = np.sum(q1 * q2, axis=-1, keepdims=True)
|
|
|
|
dot = np.where(dot < 0, -dot, dot) # handle negative dot product
|
|
|
|
dot = np.minimum(dot, 1.)
|
|
theta = np.arccos(dot) * t
|
|
|
|
q_ortho = q2 - q1 * dot
|
|
q_ortho = q_ortho / np.maximum(np.linalg.norm(q_ortho, axis=-1, keepdims=True), 1e-12)
|
|
q = q1 * np.cos(theta) + q_ortho * np.sin(theta)
|
|
return q
|
|
|
|
|
|
def slerp_rotation_matrix(R1: np.ndarray, R2: np.ndarray, t: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Spherical linear interpolation between two rotation matrices.
|
|
|
|
Args:
|
|
R1 (np.ndarray): [..., 3, 3] rotation matrix 1
|
|
R2 (np.ndarray): [..., 3, 3] rotation matrix 2
|
|
t (np.ndarray): [...] interpolation parameter in [0, 1]
|
|
|
|
Returns:
|
|
np.ndarray: [..., 3, 3] interpolated rotation matrix
|
|
"""
|
|
quat1 = matrix_to_quaternion(R1)
|
|
quat2 = matrix_to_quaternion(R2)
|
|
quat = slerp_quaternion(quat1, quat2, t)
|
|
return quaternion_to_matrix(quat)
|
|
|
|
|
|
def slerp_vector(v1: np.ndarray, v2: np.ndarray, t: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Spherical linear interpolation between two unit vectors. The vectors are assumed to be normalized.
|
|
|
|
Args:
|
|
v1 (np.ndarray): [..., d] unit vector 1
|
|
v2 (np.ndarray): [..., d] unit vector 2
|
|
t (np.ndarray): [...] interpolation parameter in [0, 1]
|
|
|
|
Returns:
|
|
np.ndarray: [..., d] interpolated unit vector
|
|
"""
|
|
dot = np.sum(v1 * v2, axis=-1, keepdims=True)
|
|
|
|
dot = np.minimum(dot, 1.)
|
|
theta = np.arccos(dot) * t
|
|
|
|
v_ortho = v2 - v1 * dot
|
|
v_ortho = v_ortho / np.maximum(np.linalg.norm(v_ortho, axis=-1, keepdims=True), 1e-12)
|
|
v = v1 * np.cos(theta) + v_ortho * np.sin(theta)
|
|
return v
|
|
|
|
|
|
def lerp(x1: np.ndarray, x2: np.ndarray, t: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Linear interpolation between two vectors.
|
|
|
|
Args:
|
|
x1 (np.ndarray): [..., d] vector 1
|
|
x2 (np.ndarray): [..., d] vector 2
|
|
t (np.ndarray): [...] interpolation parameter. [0, 1] for interpolation between x1 and x2, otherwise for extrapolation.
|
|
|
|
Returns:
|
|
np.ndarray: [..., d] interpolated vector
|
|
"""
|
|
return x1 + np.asarray(t)[..., None] * (x2 - x1)
|
|
|
|
|
|
def lerp_se3_matrix(T1: np.ndarray, T2: np.ndarray, t: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Linear interpolation between two SE(3) matrices.
|
|
|
|
Args:
|
|
T1 (np.ndarray): [..., 4, 4] SE(3) matrix 1
|
|
T2 (np.ndarray): [..., 4, 4] SE(3) matrix 2
|
|
t (np.ndarray): [...] interpolation parameter in [0, 1]
|
|
|
|
Returns:
|
|
np.ndarray: [..., 4, 4] interpolated SE(3) matrix
|
|
"""
|
|
R1 = T1[..., :3, :3]
|
|
R2 = T2[..., :3, :3]
|
|
trans1 = T1[..., :3, 3]
|
|
trans2 = T2[..., :3, 3]
|
|
R = slerp_rotation_matrix(R1, R2, t)
|
|
trans = lerp(trans1, trans2, t)
|
|
return se3_matrix(R, trans)
|
|
|
|
|
|
def piecewise_lerp(x: np.ndarray, t: np.ndarray, s: np.ndarray, extrapolation_mode: Literal['constant', 'linear'] = 'constant') -> np.ndarray:
|
|
"""
|
|
Linear spline interpolation.
|
|
|
|
### Parameters:
|
|
- `x`: np.ndarray, shape (n, d): the values of data points.
|
|
- `t`: np.ndarray, shape (n,): the times of the data points.
|
|
- `s`: np.ndarray, shape (m,): the times to be interpolated.
|
|
- `extrapolation_mode`: str, the mode of extrapolation. 'constant' means extrapolate the boundary values, 'linear' means extrapolate linearly.
|
|
|
|
### Returns:
|
|
- `y`: np.ndarray, shape (..., m, d): the interpolated values.
|
|
"""
|
|
i = np.searchsorted(t, s, side='left')
|
|
if extrapolation_mode == 'constant':
|
|
prev = np.clip(i - 1, 0, len(t) - 1)
|
|
suc = np.clip(i, 0, len(t) - 1)
|
|
elif extrapolation_mode == 'linear':
|
|
prev = np.clip(i - 1, 0, len(t) - 2)
|
|
suc = np.clip(i, 1, len(t) - 1)
|
|
else:
|
|
raise ValueError(f'Invalid extrapolation_mode: {extrapolation_mode}')
|
|
|
|
u = (s - t[prev]) / np.maximum(t[suc] - t[prev], 1e-12)
|
|
y = lerp(x[prev], x[suc], u)
|
|
|
|
return y
|
|
|
|
|
|
def piecewise_lerp_se3_matrix(T: np.ndarray, t: np.ndarray, s: np.ndarray, extrapolation_mode: Literal['constant', 'linear'] = 'constant') -> np.ndarray:
|
|
"""
|
|
Linear spline interpolation for SE(3) matrices.
|
|
|
|
### Parameters:
|
|
- `T`: np.ndarray, shape (n, 4, 4): the SE(3) matrices.
|
|
- `t`: np.ndarray, shape (n,): the times of the data points.
|
|
- `s`: np.ndarray, shape (m,): the times to be interpolated.
|
|
- `extrapolation_mode`: str, the mode of extrapolation. 'constant' means extrapolate the boundary values, 'linear' means extrapolate linearly.
|
|
|
|
### Returns:
|
|
- `T_interp`: np.ndarray, shape (..., m, 4, 4): the interpolated SE(3) matrices.
|
|
"""
|
|
i = np.searchsorted(t, s, side='left')
|
|
if extrapolation_mode == 'constant':
|
|
prev = np.clip(i - 1, 0, len(t) - 1)
|
|
suc = np.clip(i, 0, len(t) - 1)
|
|
elif extrapolation_mode == 'linear':
|
|
prev = np.clip(i - 1, 0, len(t) - 2)
|
|
suc = np.clip(i, 1, len(t) - 1)
|
|
else:
|
|
raise ValueError(f'Invalid extrapolation_mode: {extrapolation_mode}')
|
|
|
|
u = (s - t[prev]) / np.maximum(t[suc] - t[prev], 1e-12)
|
|
T = lerp_se3_matrix(T[prev], T[suc], u)
|
|
|
|
return T
|
|
|
|
|
|
def apply_transform(T: np.ndarray, x: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Apply SE(3) transformation to a point or a set of points.
|
|
|
|
### Parameters:
|
|
- `T`: np.ndarray, shape (..., 4, 4): the SE(3) matrix.
|
|
- `x`: np.ndarray, shape (..., 3): the point or a set of points to be transformed.
|
|
|
|
### Returns:
|
|
- `x_transformed`: np.ndarray, shape (..., 3): the transformed point or a set of points.
|
|
"""
|
|
x = np.asarray(x)
|
|
assert x.shape[-1] == 3
|
|
T = np.asarray(T)
|
|
assert T.shape[-2:] == (4, 4)
|
|
x_transformed = (T[..., :3, :3] @ x[..., :, None]) + T[..., :3, 3][..., None]
|
|
return x_transformed[..., 0] |