| import numpy as np |
| from numpy import ndarray |
| from typing import * |
| import itertools |
| from numbers import Number |
|
|
| from .helpers import toarray, batched |
| from ..helpers import no_warnings |
| from .utils import lite_dot, lite_norm, lite_sum |
|
|
|
|
| __all__ = [ |
| 'perspective_from_fov', |
| 'perspective_from_window', |
| 'intrinsics_from_fov', |
| 'intrinsics_from_focal_center', |
| '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', |
| 'denormalize_intrinsics', |
| 'crop_intrinsics', |
| 'pixel_to_uv', |
| 'pixel_to_ndc', |
| 'uv_to_pixel', |
| 'depth_linear_to_buffer', |
| 'depth_buffer_to_linear', |
| 'unproject_cv', |
| 'unproject_gl', |
| 'project_cv', |
| 'project_gl', |
| 'project', |
| 'unproject', |
| 'screen_coord_to_view_coord', |
| 'quaternion_to_matrix', |
| 'quaternion_multiply', |
| 'quaternion_inverse', |
| 'quaternion_normalize', |
| 'axis_angle_to_matrix', |
| 'matrix_to_quaternion', |
| 'extrinsics_to_essential', |
| 'axis_angle_to_quaternion', |
| 'euler_axis_angle_rotation', |
| 'euler_angles_to_matrix', |
| 'matrix_to_axis_angle', |
| 'matrix_to_euler_angles', |
| 'quaternion_to_axis_angle', |
| 'skew_symmetric', |
| 'rotation_matrix_from_vectors', |
| 'ray_intersection', |
| 'make_affine_matrix', |
| 'random_rotation_matrix', |
| 'lerp', |
| 'slerp', |
| 'slerp_rotation_matrix', |
| 'interpolate_se3_matrix', |
| 'piecewise_lerp', |
| 'piecewise_interpolate_se3_matrix', |
| 'transform_points', |
| 'angle_between', |
| ] |
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(_others=0) |
| def perspective_from_fov( |
| *, |
| fov_x: Optional[Union[float, ndarray]] = None, |
| fov_y: Optional[Union[float, ndarray]] = None, |
| fov_min: Optional[Union[float, ndarray]] = None, |
| fov_max: Optional[Union[float, ndarray]] = None, |
| aspect_ratio: Optional[Union[float, ndarray]] = None, |
| near: Optional[Union[float, ndarray]], |
| far: Optional[Union[float, ndarray]], |
| ) -> ndarray: |
| """ |
| Get OpenGL perspective matrix from field of view |
| |
| ## Returns |
| (ndarray): [..., 4, 4] perspective matrix |
| """ |
| if fov_max is not None: |
| fx = np.maximum(1, 1 / aspect_ratio) / np.tan(fov_max / 2) |
| fy = np.maximum(1, aspect_ratio) / np.tan(fov_max / 2) |
| elif fov_min is not None: |
| fx = np.minimum(1, 1 / aspect_ratio) / np.tan(fov_min / 2) |
| fy = np.minimum(1, aspect_ratio) / np.tan(fov_min / 2) |
| elif fov_x is not None and fov_y is not None: |
| fx = 1 / np.tan(fov_x / 2) |
| fy = 1 / np.tan(fov_y / 2) |
| elif fov_x is not None: |
| fx = 1 / np.tan(fov_x / 2) |
| fy = fx * aspect_ratio |
| elif fov_y is not None: |
| fy = 1 / np.tan(fov_y / 2) |
| fx = fy / aspect_ratio |
| perspective = np.zeros((fx.shape[0], 4, 4), dtype=fx.dtype) |
| perspective[:, 0, 0] = fx |
| perspective[:, 1, 1] = fy |
| perspective[:, 2, 2] = (near / far + 1) / (near / far - 1) |
| perspective[:, 2, 3] = 2. * near / (near / far - 1) |
| perspective[:, 3, 2] = -1. |
| return perspective |
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(_others=0) |
| def perspective_from_window( |
| left: Union[float, ndarray], |
| right: Union[float, ndarray], |
| bottom: Union[float, ndarray], |
| top: Union[float, ndarray], |
| near: Union[float, ndarray], |
| far: Union[float, ndarray] |
| ) -> ndarray: |
| """ |
| Get OpenGL perspective matrix from the window of z=-1 projection plane |
| |
| ## Returns |
| (ndarray): [..., 4, 4] perspective matrix |
| """ |
| perspective = np.zeros((left.shape[0], 4, 4), dtype=left.dtype) |
| perspective[:, 0, 0] = 2 / (right - left) |
| perspective[:, 0, 2] = (right + left) / (right - left) |
| perspective[:, 1, 1] = 2 / (top - bottom) |
| perspective[:, 1, 2] = (top + bottom) / (top - bottom) |
| perspective[:, 2, 2] = (near / far + 1) / (near / far - 1) |
| perspective[:, 2, 3] = 2. * near / (near / far - 1) |
| perspective[:, 3, 2] = -1. |
| return perspective |
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(_others=0) |
| def intrinsics_from_focal_center( |
| fx: Union[float, ndarray], |
| fy: Union[float, ndarray], |
| cx: Union[float, ndarray], |
| cy: Union[float, ndarray], |
| ) -> ndarray: |
| """ |
| Get OpenCV intrinsics matrix |
| |
| ## Returns |
| (ndarray): [..., 3, 3] OpenCV intrinsics matrix |
| """ |
| if any(isinstance(x, 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 |
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(_others=0) |
| def intrinsics_from_fov( |
| *, |
| fov_x: Optional[Union[float, ndarray]] = None, |
| fov_y: Optional[Union[float, ndarray]] = None, |
| fov_max: Optional[Union[float, ndarray]] = None, |
| fov_min: Optional[Union[float, ndarray]] = None, |
| cx: Union[float, ndarray] = 0.5, |
| cy: Union[float, ndarray] = 0.5, |
| aspect_ratio: Optional[Union[float, ndarray]] = None, |
| ) -> ndarray: |
| """ |
| Get normalized OpenCV intrinsics matrix from given field of view. |
| You can provide either fov_x, fov_y, fov_max or fov_min and aspect_ratio |
| |
| Parameters |
| ---- |
| fov_x (float | ndarray): field of view in x axis |
| fov_y (float | ndarray): field of view in y axis |
| fov_max (float | ndarray): field of view in largest dimension |
| fov_min (float | ndarray): field of view in smallest dimension |
| cx (float | ndarray): principal point x coordinate |
| cy (float | ndarray): principal point y coordinate |
| aspect_ratio (float | ndarray): aspect ratio of the image |
| |
| Returns |
| ---- |
| (ndarray): [..., 3, 3] OpenCV intrinsics matrix |
| """ |
| if fov_max is not None: |
| fx = np.maximum(1, 1 / aspect_ratio) / (2 * np.tan(fov_max / 2)) |
| fy = np.maximum(1, aspect_ratio) / (2 * np.tan(fov_max / 2)) |
| elif fov_min is not None: |
| fx = np.minimum(1, 1 / aspect_ratio) / (2 * np.tan(fov_min / 2)) |
| fy = np.minimum(1, aspect_ratio) / (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 * aspect_ratio |
| elif fov_y is not None: |
| fy = 1 / (2 * np.tan(fov_y / 2)) |
| fx = fy / aspect_ratio |
| ret = intrinsics_from_focal_center(fx, fy, cx, cy) |
| return ret |
|
|
|
|
| def focal_to_fov(focal: ndarray): |
| return 2 * np.arctan(0.5 / focal) |
|
|
|
|
| def fov_to_focal(fov: ndarray): |
| return 0.5 / np.tan(fov / 2) |
|
|
|
|
| def intrinsics_to_fov(intrinsics: ndarray) -> Tuple[ndarray, ndarray]: |
| fov_x = focal_to_fov(intrinsics[..., 0, 0]) |
| fov_y = focal_to_fov(intrinsics[..., 1, 1]) |
| return fov_x, fov_y |
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(_others=1) |
| def view_look_at( |
| eye: ndarray, |
| look_at: ndarray, |
| up: ndarray |
| ) -> ndarray: |
| """ |
| Get OpenGL view matrix looking at something |
| |
| ## Parameters |
| eye (ndarray): [..., 3] the eye position |
| look_at (ndarray): [..., 3] the position to look at |
| up (ndarray): [..., 3] head up direction (y axis in screen space). Not necessarily othogonal to view direction |
| |
| ## Returns |
| (ndarray): [..., 4, 4], view matrix |
| """ |
| z = eye - look_at |
| x = np.cross(up, z) |
| y = np.cross(z, x) |
| R = np.stack([x, y, z], axis=-2) |
| R = R / lite_norm(R, axis=-1)[..., None] |
| t = (-R @ eye[..., None]).squeeze(-1) |
| return make_affine_matrix(R, t) |
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(_others=1) |
| def extrinsics_look_at( |
| eye: ndarray, |
| look_at: ndarray, |
| up: ndarray |
| ) -> ndarray: |
| """ |
| Get OpenCV extrinsics matrix looking at something |
| |
| ## Parameters |
| eye (ndarray): [..., 3] the eye position |
| look_at (ndarray): [..., 3] the position to look at |
| up (ndarray): [..., 3] head up direction (-y axis in screen space). Not necessarily othogonal to view direction |
| |
| ## Returns |
| (ndarray): [..., 4, 4], extrinsics matrix |
| """ |
| z = look_at - eye |
| x = np.cross(-up, z) |
| y = np.cross(z, x) |
| |
| x = x / lite_norm(x, axis=-1)[..., None] |
| y = y / lite_norm(y, axis=-1)[..., None] |
| z = z / lite_norm(z, axis=-1)[..., None] |
| 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: ndarray) -> ndarray: |
| """ |
| OpenGL perspective matrix to OpenCV intrinsics |
| |
| ## Parameters |
| perspective (ndarray): [..., 4, 4] OpenGL perspective matrix |
| |
| ## Returns |
| (ndarray): shape [..., 3, 3] OpenCV intrinsics |
| """ |
| assert np.allclose(perspective[:, [0, 1, 3], 3], 0), "The matrix is not a perspective projection matrix" |
| 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: ndarray) -> Tuple[ndarray, ndarray]: |
| """ |
| Get near and far planes from OpenGL perspective matrix |
| |
| ## Parameters |
| """ |
| a, b = perspective[..., 2, 2], perspective[..., 2, 3] |
| near, far = b / (a - 1), b / (a + 1) |
| return near, far |
|
|
|
|
| @toarray(None, _others='intrinsics') |
| @batched(2, 0, 0) |
| def intrinsics_to_perspective( |
| intrinsics: ndarray, |
| near: Union[float, ndarray], |
| far: Union[float, ndarray], |
| ) -> ndarray: |
| """ |
| OpenCV intrinsics to OpenGL perspective matrix |
| |
| ## Parameters |
| intrinsics (ndarray): [..., 3, 3] OpenCV intrinsics matrix |
| near (float | ndarray): [...] near plane to clip |
| far (float | ndarray): [...] far plane to clip |
| ## Returns |
| (ndarray): [..., 4, 4] OpenGL perspective matrix |
| """ |
| perspective = np.zeros((intrinsics.shape[0], 4, 4), dtype=intrinsics.dtype) |
| perspective[..., [0, 1, 3], :3] = np.array([[2, 0, -1], [0, -2, 1], [0, 0, 1]], dtype=intrinsics.dtype) \ |
| @ intrinsics \ |
| @ np.diag(np.array([1, -1, -1], dtype=intrinsics.dtype)) |
| perspective[:, 2, 2] = (near / far + 1) / (near / far - 1) |
| perspective[:, 2, 3] = 2. * near / (near / far - 1) |
| perspective[:, 3, 2] = -1. |
| return perspective |
|
|
|
|
| def extrinsics_to_view(extrinsics: ndarray) -> ndarray: |
| """ |
| OpenCV camera extrinsics to OpenGL view matrix |
| |
| ## Parameters |
| extrinsics (ndarray): [..., 4, 4] OpenCV camera extrinsics matrix |
| |
| ## Returns |
| (ndarray): [..., 4, 4] OpenGL view matrix |
| """ |
| return extrinsics * np.array([1, -1, -1, 1], dtype=extrinsics.dtype)[:, None] |
|
|
|
|
| def view_to_extrinsics(view: ndarray) -> ndarray: |
| """ |
| OpenGL view matrix to OpenCV camera extrinsics |
| |
| ## Parameters |
| view (ndarray): [..., 4, 4] OpenGL view matrix |
| |
| ## Returns |
| (ndarray): [..., 4, 4] OpenCV camera extrinsics matrix |
| """ |
| return view * np.array([1, -1, -1, 1], dtype=view.dtype)[:, None] |
|
|
|
|
| @toarray(None, 'intrinsics') |
| @batched(2, 1) |
| def normalize_intrinsics( |
| intrinsics: ndarray, |
| size: Union[Tuple[Number, Number], ndarray], |
| pixel_convention: Literal['integer-center', 'integer-corner'] = 'integer-center', |
| ) -> ndarray: |
| """ |
| Normalize intrinsics from pixel cooridnates to uv coordinates |
| |
| ## Parameters |
| - `intrinsics` (ndarray): `(..., 3, 3)` camera intrinsics to normalize |
| - `size` (tuple | ndarray): A tuple `(height, width)` of the image size, |
| or an array of shape `(..., 2)` corresponding to the multiple image size(s) |
| - `pixel_convention`: `str`, optional `'integer-center'` or `'integer-corner'`, whether integer coordinates correspond to pixel centers or corners. Defaults to 'integer-center'. |
| - For more definitions, please refer to `pixel_coord_map()` |
| |
| ## Returns |
| `(ndarray)`: `(..., 3, 3)` normalized camera intrinsics(s) |
| """ |
| if isinstance(size, tuple): |
| size = np.array(size, dtype=intrinsics.dtype) |
| size = np.broadcast_to(size, (*intrinsics.shape[:-2], 2)) |
| height, width = size[..., 0], size[..., 1] |
| zeros = np.zeros_like(width) |
| ones = np.ones_like(width) |
| if pixel_convention == 'integer-center': |
| transform = np.stack([ |
| 1 / width, zeros, 0.5 / width, |
| zeros, 1 / height, 0.5 / height, |
| zeros, zeros, ones |
| ], axis=-1).reshape(*zeros.shape, 3, 3) |
| elif pixel_convention == 'integer-corner': |
| transform = np.stack([ |
| 1 / width, zeros, zeros, |
| zeros, 1 / height, zeros, |
| zeros, zeros, ones |
| ], axis=-1).reshape(*zeros.shape, 3, 3) |
| return transform @ intrinsics |
|
|
|
|
| @toarray(None, 'intrinsics') |
| @batched(2, 1) |
| def denormalize_intrinsics( |
| intrinsics: ndarray, |
| size: Union[Tuple[Number, Number], ndarray], |
| pixel_convention: Literal['integer-center', 'integer-corner'] = 'integer-center', |
| ) -> ndarray: |
| """ |
| Denormalize intrinsics from uv cooridnates to pixel coordinates |
| |
| ## Parameters |
| - `intrinsics` (ndarray): `(..., 3, 3)` camera intrinsics to denormalize |
| - `size` (tuple | ndarray): A tuple `(height, width)` of the image size, |
| or an array of shape `(..., 2)` corresponding to the multiple image size(s) |
| - `pixel_convention`: `str`, optional `'integer-center'` or `'integer-corner'`, whether integer coordinates correspond to pixel centers or corners. Defaults to 'integer-center'. |
| - For more definitions, please refer to `pixel_coord_map()` |
| |
| ## Returns |
| `(ndarray)`: `(..., 3, 3)` denormalized camera intrinsics in pixel coordinates |
| """ |
| if isinstance(size, tuple): |
| size = np.array(size, dtype=intrinsics.dtype) |
| size = np.broadcast_to(size, (*intrinsics.shape[:-2], 2)) |
| height, width = size[..., 0], size[..., 1] |
| zeros = np.zeros_like(width) |
| ones = np.ones_like(width) |
| if pixel_convention == 'integer-center': |
| transform = np.stack([ |
| width, zeros, -0.5 * ones, |
| zeros, height, -0.5 * ones, |
| zeros, zeros, ones |
| ], axis=-1).reshape(*zeros.shape, 3, 3) |
| elif pixel_convention == 'integer-corner': |
| transform = np.stack([ |
| width, zeros, zeros, |
| zeros, height, zeros, |
| zeros, zeros, ones |
| ], axis=-1).reshape(*zeros.shape, 3, 3) |
| return transform @ intrinsics |
|
|
|
|
| @toarray(None, _others='intrinsics') |
| @batched(2, 1, _others=0) |
| def crop_intrinsics( |
| intrinsics: ndarray, |
| size: Union[Tuple[Number, Number], ndarray], |
| cropped_top: Union[Number, ndarray], |
| cropped_left: Union[Number, ndarray], |
| cropped_height: Union[Number, ndarray], |
| cropped_width: Union[Number, ndarray], |
| ) -> ndarray: |
| """ |
| Evaluate the new intrinsics after cropping the image |
| |
| ## Parameters |
| - `intrinsics` (ndarray): (..., 3, 3) camera intrinsics(s) to crop |
| - `size` (tuple | ndarray): A tuple `(height, width)` of the image size, |
| or an array of shape `(..., 2)` corresponding to the multiple image size(s) |
| - `cropped_top` (int | ndarray): (...) top pixel index of the cropped image(s) |
| - `cropped_left` (int | ndarray): (...) left pixel index of the cropped image(s) |
| - `cropped_height` (int | ndarray): (...) height of the cropped image(s) |
| - `cropped_width` (int | ndarray): (...) width of the cropped image(s) |
| |
| ## Returns |
| (ndarray): (..., 3, 3) cropped camera intrinsics |
| """ |
| height, width = size[..., 0], size[..., 1] |
| zeros = np.zeros_like(width) |
| ones = np.ones_like(width) |
| transform = np.stack([ |
| width / cropped_width, zeros, -cropped_left / cropped_width, |
| zeros, height / cropped_height, -cropped_top / cropped_height, |
| zeros, zeros, ones |
| ]).reshape(*zeros.shape, 3, 3) |
| return transform @ intrinsics |
|
|
|
|
| def pixel_to_uv( |
| pixel: ndarray, |
| size: Union[Tuple[Number, Number], ndarray], |
| pixel_convention: Literal['integer-center', 'integer-corner'] = 'integer-center', |
| ) -> ndarray: |
| """ |
| Convert pixel space coordiantes to UV space coordinates. |
| |
| ## Parameters |
| - `pixel` (ndarray): `(..., 2)` pixel coordinrates |
| - `size` (tuple | ndarray): A tuple `(height, width)` of the image size, |
| or an array of shape `(..., 2)` corresponding to the multiple image size(s) |
| - `pixel_convention`: `str`, optional `'integer-center'` or `'integer-corner'`, whether integer coordinates correspond to pixel centers or corners. Defaults to 'integer-center'. |
| - For more definitions, please refer to `pixel_coord_map()` |
| |
| ## Returns |
| (ndarray): `(..., 2)` uv coordinrates |
| """ |
| if not np.issubdtype(pixel.dtype, np.floating): |
| pixel = pixel.astype(np.float32) |
| if pixel_convention == 'integer-center': |
| pixel = pixel + 0.5 |
| uv = pixel / np.flip(size, axis=-1) |
| return uv |
|
|
|
|
| @toarray(size='uv') |
| def uv_to_pixel( |
| uv: ndarray, |
| size: Union[Tuple[Number, Number], ndarray], |
| pixel_convention: Literal['integer-center', 'integer-corner'] = 'integer-center', |
| ) -> ndarray: |
| """ |
| Convert UV space coordinates to pixel space coordinates. |
| |
| ## Parameters |
| - `uv` (ndarray): `(..., 2)` uv coordinrates. |
| - `size` (tuple | ndarray): A tuple `(height, width)` of the image size, |
| or an array of shape `(..., 2)` corresponding to the multiple image size(s) |
| - `pixel_convention`: `str`, optional `'integer-center'` or `'integer-corner'`, whether integer coordinates correspond to pixel centers or corners. Defaults to 'integer-center'. |
| - For more definitions, please refer to `pixel_coord_map()` |
| |
| ## Returns |
| (ndarray): `(..., 2)` pixel coordinrates |
| """ |
| pixel = uv * np.flip(size, axis=-1) |
| if pixel_convention == 'integer-center': |
| pixel = pixel - 0.5 |
| return pixel |
|
|
|
|
| @toarray(size=np.float32) |
| def pixel_to_ndc( |
| pixel: ndarray, |
| size: Union[Tuple[Number, Number], ndarray], |
| pixel_convention: Literal['integer-center', 'integer-corner'] = 'integer-center', |
| ) -> ndarray: |
| """ |
| Convert pixel coordinates to NDC (Normalized Device Coordinates). |
| |
| ## Parameters |
| - `pixel` (ndarray): `(..., 2)` pixel coordinrates. |
| - `size` (tuple | ndarray): A tuple `(height, width)` of the image size, |
| or an array of shape `(..., 2)` corresponding to the multiple image size(s) |
| - `pixel_convention`: `str`, optional `'integer-center'` or `'integer-corner'`, whether integer coordinates represent pixel centers or corners. Defaults to 'integer-center'. |
| - For more definitions, please refer to `pixel_coord_map()` |
| |
| ## Returns |
| (ndarray): `(..., 2)` ndc coordinrates, the range is (-1, 1) |
| """ |
| if not np.issubdtype(pixel.dtype, np.floating): |
| pixel = pixel.astype(np.float32) |
| dtype = pixel.dtype |
| if pixel_convention == 'integer-center': |
| pixel = pixel + 0.5 |
| ndc = pixel / (np.flip(size, axis=-1) * np.array([2, -2], dtype=dtype)) \ |
| + np.array([-1, 1], dtype=dtype) |
| return ndc |
|
|
|
|
| @batched(0, 0, 0) |
| def depth_linear_to_buffer( |
| depth: ndarray, |
| near: Union[float, ndarray], |
| far: Union[float, ndarray] |
| ) -> ndarray: |
| """ |
| Project linear depth to depth value in screen space |
| |
| ## Parameters |
| depth (ndarray): [...] depth value |
| near (float | ndarray): [...] near plane to clip |
| far (float | ndarray): [...] far plane to clip |
| |
| ## Returns |
| (ndarray): [..., 1] depth value in screen space, value ranging in [0, 1] |
| """ |
| return (1 - near / depth) / (1 - near / far) |
|
|
|
|
| @batched(0, 0, 0) |
| def depth_buffer_to_linear( |
| depth_buffer: ndarray, |
| near: Union[float, ndarray], |
| far: Union[float, ndarray] |
| ) -> ndarray: |
| """ |
| OpenGL depth buffer to linear depth |
| |
| ## Parameters |
| depth_buffer (ndarray): [...] depth value |
| near (float | ndarray): [...] near plane to clip |
| far (float | ndarray): [...] far plane to clip |
| |
| ## Returns |
| (ndarray): [..., 1] linear depth |
| """ |
| return near / (1 - (1 - near / far) * depth_buffer) |
|
|
|
|
| def project_gl( |
| points: ndarray, |
| projection: ndarray, |
| view: ndarray = None, |
| ) -> Tuple[ndarray, ndarray]: |
| """ |
| Project 3D points to 2D following the OpenGL convention (except for row major matrices) |
| |
| ## Parameters |
| points (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 |
| view (ndarray): [..., 4, 4] view matrix |
| projection (ndarray): [..., 4, 4] projection matrix |
| |
| ## Returns |
| scr_coord (ndarray): [..., N, 2] OpenGL screen space XY coordinates, value ranging in [0, 1]. |
| The origin (0., 0.) is corresponding to the left & bottom |
| linear_depth (ndarray): [..., N] linear depth |
| """ |
| if points.shape[-1] == 3: |
| points = np.concatenate([points, np.ones((*points.shape[:-1], 1), dtype=points.dtype)], axis=-1) |
| transform = projection @ view if view is not None else projection |
| clip_coord = points @ transform.swapaxes(-2, -1) |
| ndc_coord = clip_coord[..., :3] / clip_coord[..., 3:] |
| scr_coord = ndc_coord * 0.5 + 0.5 |
| linear_depth = clip_coord[..., 3] |
| return scr_coord[..., :2], linear_depth |
|
|
|
|
| @no_warnings() |
| def project_cv( |
| points: ndarray, |
| intrinsics: ndarray, |
| extrinsics: Optional[ndarray] = None, |
| ) -> Tuple[ndarray, ndarray]: |
| """ |
| Project 3D points to 2D following the OpenCV convention |
| |
| ## Parameters |
| points (ndarray): [..., N, 3] |
| extrinsics (ndarray): [..., 4, 4] extrinsics matrix |
| intrinsics (ndarray): [..., 3, 3] intrinsics matrix |
| |
| ## Returns |
| uv_coord (ndarray): [..., N, 2] uv coordinates, value ranging in [0, 1]. |
| The origin (0., 0.) is corresponding to the left & top |
| linear_depth (ndarray): [..., N] linear depth |
| """ |
| assert intrinsics is not None, "intrinsics matrix is required" |
| points = np.concatenate([points, np.ones((*points.shape[:-1], 1), dtype=points.dtype)], axis=-1) |
| intrinsics = np.block([ |
| [intrinsics, np.zeros((*intrinsics.shape[:-2], 3, 1), dtype=intrinsics.dtype)], |
| [np.broadcast_to(np.array([0, 0, 0, 1], dtype=intrinsics.dtype), (*intrinsics.shape[:-2], 1, 4))] |
| ]) |
| transform = intrinsics @ extrinsics if extrinsics is not None else intrinsics |
| points = points @ transform.swapaxes(-2, -1) |
| uv_coord = points[..., :2] / points[..., 2:3] |
| linear_depth = points[..., 2] |
| return uv_coord, linear_depth |
|
|
|
|
| def unproject_gl( |
| uv: ndarray, |
| depth: ndarray, |
| projection: ndarray, |
| view: Optional[ndarray] = None, |
| ) -> ndarray: |
| """ |
| Unproject screen space coordinates to 3D view space following the OpenGL convention (except for row major matrices) |
| |
| ## Parameters |
| uv (ndarray): (..., N, 2) screen space XY coordinates, value ranging in [0, 1]. |
| The origin (0., 0.) is corresponding to the left & bottom |
| depth (ndarray): (..., N) linear depth values |
| projection (ndarray): (..., 4, 4) projection matrix |
| view (ndarray): (..., 4, 4) view matrix |
| |
| ## Returns |
| points (ndarray): (..., N, 3) 3d points |
| """ |
| ndc_xy = uv * 2 - 1 |
| view_z = -depth |
| clip_xy = np.linalg.inv(projection[..., :2, :2] - ndc_xy[..., :, None] * projection[..., 3:, :2]) \ |
| @ ((ndc_xy[..., :, None] * projection[..., 3:, 2:] - projection[..., :2, 2:]) \ |
| @ np.concatenate([view_z[..., None, None], np.ones_like(view_z[..., None, None])], axis=-2)) |
| points = np.concatenate([clip_xy.squeeze(-1), view_z[..., None], np.ones_like(view_z)[..., None]], axis=-1) |
| if view is not None: |
| points = points @ np.linalg.inv(view).swapaxes(-2, -1) |
| return points[..., :3] |
|
|
|
|
| @batched(2, 2) |
| def screen_coord_to_view_coord(screen_coord: ndarray, projection: ndarray) -> ndarray: |
| """ |
| Unproject screen space coordinates to 3D view space following the OpenGL convention (except for row major matrices) |
| |
| ## Parameters |
| screen_coord (ndarray): (..., N, 3) screen space XYZ coordinates, value ranging in [0, 1] |
| The origin (0., 0.) is corresponding to the left & bottom |
| projection (ndarray): (..., 4, 4) projection matrix |
| |
| ## Returns |
| points (ndarray): [..., N, 3] 3d points in view space |
| """ |
| assert projection is not None, "projection matrix is required" |
| ndc_xy = screen_coord * 2 - 1 |
| clip_coord = np.concatenate([ndc_xy, np.ones_like(ndc_xy[..., :1])], axis=-1) |
| points = clip_coord @ np.linalg.inv(projection).swapaxes(-1, -2) |
| points = points[..., :3] / points[..., 3:] |
| return points |
|
|
|
|
| @batched(2, 1, 2, 2) |
| def unproject_cv( |
| uv: ndarray, |
| depth: ndarray, |
| intrinsics: ndarray, |
| extrinsics: ndarray = None, |
| ) -> ndarray: |
| """ |
| Unproject uv coordinates to 3D view space following the OpenCV convention |
| |
| ## Parameters |
| uv (ndarray): [..., N, 2] uv coordinates, value ranging in [0, 1]. |
| The origin (0., 0.) is corresponding to the left & top |
| depth (ndarray): [..., N] depth value |
| extrinsics (ndarray): [..., 4, 4] extrinsics matrix |
| intrinsics (ndarray): [..., 3, 3] intrinsics matrix |
| |
| ## Returns |
| points (ndarray): [..., N, 3] 3d points |
| """ |
| intrinsics = np.block([ |
| [intrinsics, np.zeros((*intrinsics.shape[:-2], 3, 1), dtype=intrinsics.dtype)], |
| [np.broadcast_to(np.array([0, 0, 0, 1], dtype=intrinsics.dtype), (*intrinsics.shape[:-2], 1, 4))] |
| ]) |
| transform = intrinsics @ extrinsics if extrinsics is not None else intrinsics |
| points = np.concatenate([uv, np.ones((*uv.shape[:-1], 1), dtype=uv.dtype)], axis=-1) * depth[..., None] |
| points = np.concatenate([points, np.ones((*points.shape[:-1], 1), dtype=uv.dtype)], axis=-1) |
| points = points @ np.linalg.inv(transform).swapaxes(-2, -1) |
| points = points[..., :3] |
| return points |
|
|
|
|
|
|
| def project( |
| points: ndarray, |
| *, |
| intrinsics: Optional[ndarray] = None, |
| extrinsics: Optional[ndarray] = None, |
| view: Optional[ndarray] = None, |
| projection: Optional[ndarray] = None |
| ) -> Tuple[ndarray, ndarray]: |
| """ |
| Calculate projection. |
| - For OpenCV convention, use `intrinsics` and `extrinsics` matrices. |
| - For OpenGL convention, use `view` and `projection` matrices. |
| |
| ## Parameters |
| |
| - `points`: (..., N, 3) 3D world-space points |
| - `intrinsics`: (..., 3, 3) intrinsics matrix |
| - `extrinsics`: (..., 4, 4) extrinsics matrix |
| - `view`: (..., 4, 4) view matrix |
| - `projection`: (..., 4, 4) projection matrix |
| |
| ## Returns |
| |
| - `uv`: (..., N, 2) 2D coordinates. |
| - For OpenCV convention, it is the normalized image coordinate where (0, 0) is the top left corner. |
| - For OpenGL convention, it is the screen space XY coordinate where (0, 0) is the bottom left corner. |
| - `depth`: (..., N) linear depth values, where `depth > 0` is visible. |
| - For OpenCV convention, it is the Z coordinate in camera space. |
| - For OpenGL convention, it is the -Z coordinate in camera space. |
| """ |
| assert (intrinsics is not None or extrinsics is not None) ^ (view is not None or projection is not None), \ |
| "Either camera intrinsics (and extrinsics) or projection (and view) matrices must be provided." |
| |
| if intrinsics is not None: |
| return project_cv(points, intrinsics, extrinsics) |
| elif projection is not None: |
| return project_gl(points, projection, view) |
| else: |
| raise ValueError("Invalid combination of input parameters.") |
|
|
|
|
| def unproject( |
| uv: ndarray, |
| depth: Optional[ndarray], |
| *, |
| intrinsics: Optional[ndarray] = None, |
| extrinsics: Optional[ndarray] = None, |
| projection: Optional[ndarray] = None, |
| view: Optional[ndarray] = None, |
| ) -> ndarray: |
| """ |
| Calculate inverse projection. |
| - For OpenCV convention, use `intrinsics` and `extrinsics` matrices. |
| - For OpenGL convention, use `view` and `projection` matrices. |
| |
| ## Parameters |
| |
| - `uv`: (..., N, 2) 2D coordinates. |
| - For OpenCV convention, it is the normalized image coordinate where (0, 0) is the top left corner. |
| - For OpenGL convention, it is the screen space XY coordinate where (0, 0) is the bottom left corner. |
| - `depth`: (..., N) linear depth values, where `depth > 0` is visible. |
| - For OpenCV convention, it is the Z coordinate in camera space. |
| - For OpenGL convention, it is the -Z coordinate in camera space. |
| - `intrinsics`: (..., 3, 3) intrinsics matrix |
| - `extrinsics`: (..., 4, 4) extrinsics matrix |
| - `view`: (..., 4, 4) view matrix |
| - `projection`: (..., 4, 4) projection matrix |
| |
| ## Returns |
| |
| - `points`: (..., N, 3) 3D world-space points |
| """ |
| assert (intrinsics is not None or extrinsics is not None) ^ (view is not None or projection is not None), \ |
| "Either camera intrinsics (and extrinsics) or projection (and view) matrices must be provided." |
|
|
| if intrinsics is not None: |
| return unproject_cv(uv, depth, intrinsics, extrinsics) |
| elif projection is not None: |
| return unproject_gl(uv, depth, projection, view) |
| else: |
| raise ValueError("Invalid combination of input parameters.") |
|
|
|
|
| def quaternion_to_matrix(quaternion: ndarray) -> ndarray: |
| """Converts a batch of quaternions (w, x, y, z) to rotation matrices |
| |
| ## Parameters |
| quaternion (ndarray): shape (..., 4), the quaternions to convert |
| |
| ## Returns |
| ndarray: shape (..., 3, 3), the rotation matrices corresponding to the given quaternions |
| """ |
| assert quaternion.shape[-1] == 4 |
| quaternion = quaternion / np.maximum(lite_norm(quaternion, axis=-1), np.finfo(quaternion.dtype).tiny)[..., None] |
| 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 * lite_sum(np.square(xyz), 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: ndarray) -> ndarray: |
| """Convert 3x3 rotation matrix to quaternion (w, x, y, z) |
| |
| ## Parameters |
| rot_mat (ndarray): shape (..., 3, 3), the rotation matrices to convert |
| |
| ## Returns |
| ndarray: shape (..., 4), the quaternions corresponding to the given rotation matrices |
| """ |
| |
| 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.maximum(lite_norm(quat, axis=-1), np.finfo(quat.dtype).tiny)[..., None] |
| return quat |
|
|
|
|
| def quaternion_multiply(q1: ndarray, q2: ndarray) -> ndarray: |
| """Multiplies two quaternions (w, x, y, z) |
| |
| Parameters |
| ---- |
| q1 (ndarray): shape (..., 4), the first quaternion |
| q2 (ndarray): shape (..., 4), the second quaternion |
| |
| Returns |
| ---- |
| ndarray: shape (..., 4), the product of the two quaternions |
| """ |
| w1, v1 = q1[..., :1], q1[..., 1:] |
| w2, v2 = q2[..., :1], q2[..., 1:] |
| res_w = w1 * w2 - lite_dot(v1, v2)[..., None] |
| res_v = w1 * v2 + w2 * v1 + np.cross(v1, v2, axis=-1) |
| res_q = np.concatenate([res_w, res_v], axis=-1) |
| return res_q |
|
|
|
|
| def quaternion_normalize(quaternion: ndarray) -> ndarray: |
| """Normalize quaternions (w, x, y, z) to unit length and positive w component |
| |
| Parameters |
| ---- |
| quaternion (ndarray): shape (..., 4), the quaternions to normalize |
| |
| Returns |
| ---- |
| ndarray: shape (..., 4), the normalized quaternions with unit length and positive w component |
| """ |
| w_sign = np.where(quaternion[..., 0] >= 0, 1.0, -1.0).astype(quaternion.dtype) |
| quat = quaternion * w_sign[..., None] |
| quat = quat / np.maximum(lite_norm(quat, axis=-1)[..., None], np.finfo(quat.dtype).tiny) |
| return quat |
|
|
|
|
| def quaternion_inverse(quaternion: ndarray) -> ndarray: |
| """Calculate the inverse of a batch of quaternions (w, x, y, z) |
| |
| Parameters |
| ---- |
| quaternion (ndarray): shape (..., 4), the quaternions to invert |
| |
| Returns |
| ---- |
| ndarray: shape (..., 4) |
| """ |
| w, v = quaternion[..., 0:1], quaternion[..., 1:] |
| inv_quat = np.concatenate([w, -v], axis=-1) |
| return inv_quat |
|
|
|
|
| def quaternion_to_axis_angle(quaternion: ndarray) -> ndarray: |
| """Convert a batch of quaternions (w, x, y, z) to axis-angle representation (rotation vector) |
| |
| ## Parameters |
| quaternion (ndarray): shape (..., 4), the quaternions to convert |
| |
| ## Returns |
| ndarray: shape (..., 3), the axis-angle vectors corresponding to the given quaternions |
| """ |
| assert quaternion.shape[-1] == 4 |
| norm = lite_norm(quaternion[..., 1:], axis=-1) |
| axis = quaternion[..., 1:] / np.maximum(norm, np.finfo(quaternion.dtype).tiny)[..., None] |
| angle = 2 * np.atan2(norm, quaternion[..., 0:1]) |
| return angle * axis |
|
|
|
|
| def matrix_to_axis_angle(rot_mat: ndarray) -> ndarray: |
| """Convert a batch of 3x3 rotation matrices to axis-angle representation (rotation vector) |
| |
| ## Parameters |
| rot_mat (ndarray): shape (..., 3, 3), the rotation matrices to convert |
| |
| ## Returns |
| ndarray: shape (..., 3), the axis-angle vectors corresponding to the given rotation matrices |
| """ |
| quat = matrix_to_quaternion(rot_mat) |
| axis_angle = quaternion_to_axis_angle(quat) |
| return axis_angle |
|
|
|
|
| def extrinsics_to_essential(extrinsics: ndarray): |
| """ |
| extrinsics matrix `[[R, t] [0, 0, 0, 1]]` such that `x' = R (x - t)` to essential matrix such that `x' E x = 0` |
| |
| ## Parameters |
| 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: ndarray) -> 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. |
| |
| ## Parameters |
| axis: Axis label "X" or "Y or "Z". |
| angle: any shape ndarray of Euler angles in radians |
| |
| ## Returns |
| Rotation matrices as ndarray 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: ndarray, convention: str = 'XYZ') -> ndarray: |
| """ |
| Convert rotations given as Euler angles in radians to rotation matrices. |
| |
| ## Parameters |
| 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: 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: ndarray, v2: ndarray): |
| "Rotation matrix that rotates v1 to v2" |
| I = np.eye(3, dtype=v1.dtype) |
| v1 = v1 / lite_norm(v1, axis=-1)[..., None] |
| v2 = v2 / lite_norm(v2, axis=-1)[..., None] |
| v = np.cross(v1, v2, axis=-1) |
| c = lite_dot(v1, v2, axis=-1) |
| K = skew_symmetric(v) |
| R = I + K + (1 / (1 + c)).astype(v1.dtype)[None, None] * (K @ K) |
| return R |
|
|
|
|
| def axis_angle_to_matrix(axis_angle: ndarray) -> ndarray: |
| """Convert axis-angle representation (rotation vector) to rotation matrix, whose direction is the axis of rotation and length is the angle of rotation |
| |
| ## Parameters |
| axis_angle (ndarray): shape (..., 3), axis-angle vcetors |
| |
| ## Returns |
| ndarray: shape (..., 3, 3) The rotation matrices for the given axis-angle parameters |
| """ |
| batch_shape = axis_angle.shape[:-1] |
| dtype = axis_angle.dtype |
|
|
| angle = lite_norm(axis_angle, axis=-1)[..., None] |
| axis = axis_angle / np.maximum(angle, np.finfo(dtype).tiny) |
|
|
| 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 axis_angle_to_quaternion(axis_angle: ndarray) -> ndarray: |
| """Convert axis-angle representation (rotation vector) to quaternion (w, x, y, z) |
| |
| ## Parameters |
| axis_angle (ndarray): shape (..., 3), axis-angle vcetors |
| |
| ## Returns |
| ndarray: shape (..., 4) The quaternions for the given axis-angle parameters |
| """ |
| angle = lite_norm(axis_angle, axis=-1)[..., None] |
| axis = axis_angle / np.maximum(angle, np.finfo(axis_angle.dtype).tiny) |
| quat = np.concatenate([np.cos(angle / 2), np.sin(angle / 2) * axis], axis=-1) |
| return quat |
|
|
|
|
| def _angle_from_tan( |
| axis: str, other_axis: str, data: ndarray, horizontal: bool, tait_bryan: bool |
| ) -> ndarray: |
| """ |
| Extract the first or third Euler angle from the two members of |
| the matrix which are positive constant times its sine and cosine. |
| |
| ## Parameters |
| axis: Axis label "X" or "Y or "Z" for the angle we are finding. |
| other_axis: Axis label "X" or "Y or "Z" for the middle axis in the |
| convention. |
| data: Rotation matrices as ndarray of shape (..., 3, 3). |
| horizontal: Whether we are looking for the angle for the third axis, |
| which means the relevant entries are in the same row of the |
| rotation matrix. If not, they are in the same column. |
| tait_bryan: Whether the first and third axes in the convention differ. |
| |
| ## Returns |
| Euler Angles in radians for each matrix in data as a ndarray |
| of shape (...). |
| """ |
|
|
| i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis] |
| if horizontal: |
| i2, i1 = i1, i2 |
| even = (axis + other_axis) in ["XY", "YZ", "ZX"] |
| if horizontal == even: |
| return np.atan2(data[..., i1], data[..., i2]) |
| if tait_bryan: |
| return np.atan2(-data[..., i2], data[..., i1]) |
| return np.atan2(data[..., i2], -data[..., i1]) |
|
|
|
|
| def matrix_to_euler_angles(matrix: ndarray, convention: str) -> ndarray: |
| """ |
| Convert rotations given as rotation matrices to Euler angles in radians. |
| NOTE: The composition order eg. `XYZ` means `Rz * Ry * Rx` (like blender), instead of `Rx * Ry * Rz` (like pytorch3d) |
| |
| ## Parameters |
| matrix: Rotation matrices as ndarray of shape (..., 3, 3). |
| convention: Convention string of three uppercase letters. |
| |
| ## Returns |
| Euler angles in radians as ndarray of shape (..., 3), in the order of XYZ (like blender), instead of convention (like pytorch3d) |
| """ |
| if not all(c in 'XYZ' for c in convention) or not all(c in convention for c in 'XYZ'): |
| raise ValueError(f"Invalid convention {convention}.") |
| if not matrix.shape[-2:] == (3, 3): |
| raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") |
| |
| i0 = 'XYZ'.index(convention[0]) |
| i2 = 'XYZ'.index(convention[2]) |
| tait_bryan = i0 != i2 |
| if tait_bryan: |
| central_angle = np.asin(matrix[..., i2, i0] * (-1.0 if i2 - i0 in [-1, 2] else 1.0)) |
| else: |
| central_angle = np.acos(matrix[..., i2, i2]) |
|
|
| |
| o = [ |
| _angle_from_tan( |
| convention[0], convention[1], matrix[..., i2, :], True, tait_bryan |
| ), |
| central_angle, |
| _angle_from_tan( |
| convention[2], convention[1], matrix[..., i0], False, tait_bryan |
| ), |
| ] |
| return np.stack([o[convention.index(c)] for c in 'XYZ'], -1) |
|
|
|
|
| def random_rotation_matrix(*size: int, dtype=np.float32) -> ndarray: |
| """ |
| Generate random 3D rotation matrix. |
| |
| ## Parameters |
| dtype: The data type of the output rotation matrix. |
| |
| ## Returns |
| ndarray: `(*size, 3, 3)` random rotation matrix. |
| """ |
| if len(size) == 1 and isinstance(size[0], (tuple, list)): |
| size = size[0] |
| rand_quat = np.random.randn(*size, 4).astype(dtype) |
| return quaternion_to_matrix(rand_quat) |
|
|
|
|
| def ray_intersection(p1: ndarray, d1: ndarray, p2: ndarray, d2: ndarray): |
| """ |
| Compute the intersection/closest point of two D-dimensional rays |
| If the rays are intersecting, the closest point is the intersection point. |
| |
| ## Parameters |
| p1 (ndarray): (..., D) origin of ray 1 |
| d1 (ndarray): (..., D) direction of ray 1 |
| p2 (ndarray): (..., D) origin of ray 2 |
| d2 (ndarray): (..., D) direction of ray 2 |
| |
| ## Returns |
| (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) |
| p = np.stack([p1, p2], axis=-2) |
| A = np.concatenate([ |
| (np.eye(dim, dtype=dtype) * np.ones((*p.shape[:-2], 2, 1, 1))).reshape(*d.shape[:-2], 2 * dim, dim), |
| -(np.eye(2, dtype=dtype)[..., None] * d[..., None, :]).swapaxes(-2, -1).reshape(*d.shape[:-2], 2 * dim, 2) |
| ], axis=-1) |
| b = p.reshape(*p.shape[:-2], 2 * dim) |
| 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]) |
|
|
|
|
| @batched(2, 1) |
| def make_affine_matrix(M: ndarray, t: ndarray) -> ndarray: |
| """ |
| Make an affine transformation matrix from a linear matrix and a translation vector. |
| |
| ## Parameters |
| M (ndarray): [..., D, D] linear matrix (rotation, scaling or general deformation) |
| t (ndarray): [..., D] translation vector |
| |
| ## Returns |
| ndarray: [..., D + 1, D + 1] affine transformation matrix |
| """ |
| x = np.block([ |
| [M, t[..., None]], |
| [np.zeros((*M.shape[:-2], 1, M.shape[-1]), dtype=M.dtype), np.ones((*M.shape[:-2], 1, 1), dtype=M.dtype)] |
| ]) |
| return x |
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(1, 1, 1) |
| def lerp(x1: ndarray, x2: ndarray, t: ndarray) -> ndarray: |
| """ |
| Linear interpolation between two vectors. |
| |
| ## Parameters |
| x1 (ndarray): [..., D] vector 1 |
| x2 (ndarray): [..., D] vector 2 |
| t (ndarray): [..., N] interpolation parameter. [0, 1] for interpolation between x1 and x2, otherwise for extrapolation. |
| |
| ## Returns |
| ndarray: [..., N, D] interpolated vector |
| """ |
| return x1[..., None, :] + t[..., None] * (x2 - x1)[..., None, :] |
|
|
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(1, 1, 1) |
| def slerp(v1: ndarray, v2: ndarray, t: ndarray) -> ndarray: |
| """ |
| Spherical linear interpolation between two (unit) vectors. |
| |
| ## Parameters |
| - `v1` (ndarray): `(..., D)` (unit) vector 1 |
| - `v2` (ndarray): `(..., D)` (unit) vector 2 |
| - `t` (ndarray): `(..., N)` interpolation parameter in [0, 1] |
| |
| ## Returns |
| ndarray: `(..., N, D)` interpolated unit vector |
| """ |
| v1 = v1 / np.maximum(lite_norm(v1, axis=-1), np.finfo(v1.dtype).tiny)[..., None] |
| v2 = v2 / np.maximum(lite_norm(v2, axis=-1), np.finfo(v2.dtype).tiny)[..., None] |
| cos = lite_dot(v1, v2, axis=-1) |
| v_ortho1 = v2 - v1 * cos[..., None] |
| v_ortho2 = v1 - v2 * cos[..., None] |
| sin = np.minimum(lite_norm(v_ortho1, axis=-1), lite_norm(v_ortho2, axis=-1)) |
| theta = np.atan2(sin, cos)[..., None] * t |
| v_ortho1 = v_ortho1 / np.maximum(lite_norm(v_ortho1, axis=-1), np.finfo(v1.dtype).tiny)[..., None] |
| v = v1[..., None, :] * np.cos(theta)[..., None] + v_ortho1[..., None, :] * np.sin(theta)[..., None] |
| return v |
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(2, 2, 1) |
| def slerp_rotation_matrix(R1: ndarray, R2: ndarray, t: ndarray) -> ndarray: |
| """ |
| Spherical linear interpolation between two rotation matrices. |
| |
| ## Parameters |
| - `R1` (ndarray): [..., 3, 3] rotation matrix 1 |
| - `R2` (ndarray): [..., 3, 3] rotation matrix 2 |
| - `t` (ndarray): [..., N] interpolation parameter in [0, 1] |
| |
| ## Returns |
| ndarray: [...,N, 3, 3] interpolated rotation matrix |
| """ |
| quat1 = matrix_to_quaternion(R1) |
| quat2 = matrix_to_quaternion(R2) |
| quat = slerp(quat1, quat2, t) |
| return quaternion_to_matrix(quat) |
|
|
|
|
| @toarray(_others=np.float32) |
| @batched(2, 2, 1) |
| def interpolate_se3_matrix(T1: ndarray, T2: ndarray, t: ndarray) -> ndarray: |
| """ |
| Linear interpolation between two SE(3) matrices. |
| |
| ## Parameters |
| - `T1` (ndarray): [..., 4, 4] SE(3) matrix 1 |
| - `T2` (ndarray): [..., 4, 4] SE(3) matrix 2 |
| - `t` (ndarray): [..., N] interpolation parameter in [0, 1] |
| |
| ## Returns |
| ndarray: [..., N, 4, 4] interpolated SE(3) matrix |
| """ |
| assert T1.shape[-2:] == (4, 4) and T2.shape[-2:] == (4, 4) |
| rot = slerp_rotation_matrix(T1[..., :3, :3], T2[..., :3, :3], t) |
| pos = lerp(T1[..., :3, 3], T2[..., :3, 3], t) |
| return make_affine_matrix(rot, pos) |
|
|
|
|
| def piecewise_lerp(x: ndarray, t: ndarray, s: ndarray, extrapolation_mode: Literal['constant', 'linear'] = 'constant') -> ndarray: |
| """ |
| Linear spline interpolation. |
| |
| ## Parameters |
| - `x`: ndarray, shape (n, d): the values of data points. |
| - `t`: ndarray, shape (n,): the times of the data points. |
| - `s`: 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`: 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_interpolate_se3_matrix(T: ndarray, t: ndarray, s: ndarray, extrapolation_mode: Literal['constant', 'linear'] = 'constant') -> ndarray: |
| """ |
| Linear spline interpolation for SE(3) matrices. |
| |
| ## Parameters |
| - `T`: ndarray, shape (n, 4, 4): the SE(3) matrices. |
| - `t`: ndarray, shape (n,): the times of the data points. |
| - `s`: 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`: 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 = interpolate_se3_matrix(T[prev], T[suc], u) |
|
|
| return T |
|
|
|
|
| def transform_points(x: ndarray, *Ts: ndarray) -> ndarray: |
| """ |
| Apply transformation(s) to a point or a set of points. |
| It is like `(Tn @ ... @ T2 @ T1 @ x[:, None]).squeeze(0)`, but: |
| 1. Automatically handle the homogeneous coordinate; |
| - x will be padded with homogeneous coordinate 1. |
| - Each T will be padded by identity matrix to match the dimension. |
| 2. Using efficient contraction path when array sizes are large, based on `einsum`. |
| |
| ## Parameters |
| - `x`: ndarray, shape `(..., D)`: the points to be transformed. |
| - `Ts`: ndarray, shape `(..., D1, D2)`: the affine transformation matrix (matrices) |
| If more than one transformation is given, they will be applied in corresponding order. |
| ## Returns |
| - `y`: ndarray, shape `(..., D)`: the transformed point or a set of points. |
| |
| ## Example Usage |
| |
| - Just linear transformation |
| |
| ``` |
| y = transform(x_3, mat_3x3) |
| ``` |
| |
| - Affine transformation |
| |
| ``` |
| y = transform(x_3, mat_3x4) |
| ``` |
| |
| - Chain multiple transformations |
| |
| ``` |
| y = transform(x_3, T1_4x4, T2_3x4, T3_3x4) |
| ``` |
| """ |
| input_dim = x.shape[-1] |
| pad_dim = max(max(max(T.shape[-2:]) for T in Ts), x.shape[-1]) |
| x = np.concatenate([x, np.ones((*x.shape[:-1], pad_dim - x.shape[-1]), dtype=x.dtype)], axis=-1) |
| I = np.eye(pad_dim, dtype=x.dtype) |
| Ts = [ |
| np.concatenate([ |
| np.concatenate([T, np.broadcast_to(I[:T.shape[-2], T.shape[-1]:], (*T.shape[:-2], T.shape[-2], pad_dim - T.shape[-1]))], axis=-1), |
| np.broadcast_to(I[T.shape[-2]:, :], (*T.shape[:-2], pad_dim - T.shape[-2], pad_dim)) |
| ], axis=-2) |
| for T in Ts |
| ] |
| total_numel = sum(t.size for t in Ts) + x.size |
| if total_numel > 1000: |
| |
| operands = [*reversed(Ts), x[..., None]] |
| offset = len(operands) + 1 |
| batch_shape = np.broadcast_shapes(*(m.shape[:-2] for m in operands)) |
| batch_subscripts = tuple(range(offset, offset + len(batch_shape))) |
| |
| subscripts, squeezed_operands = [], [] |
| for i, m in enumerate(operands): |
| squeezable = tuple(b_m == 1 and b > 1 for b_m, b in zip(m.shape[:-2], batch_shape[len(batch_shape) - (m.ndim - 2):])) |
| squeezed_operands.append( |
| m.squeeze(axis=tuple(j for j, s in enumerate(squeezable) if s)) |
| ) |
| subscripts.append( |
| (*tuple(j for j, s in zip(batch_subscripts[len(batch_shape) - (m.ndim - 2):], squeezable) if not s), i, i + 1) |
| ) |
| y = np.einsum( |
| *itertools.chain(*zip(squeezed_operands, subscripts)), |
| (*range(offset, offset + len(batch_shape)), 0, len(operands)), |
| optimize="optimal" |
| ) |
| y = y.squeeze(-1) |
| else: |
| y = x[..., None] |
| for T in Ts: |
| y = T @ y |
| y = y.squeeze(-1) |
| return y[..., :input_dim] |
| |
|
|
| def angle_between(v1: ndarray, v2: ndarray): |
| """ |
| Calculate the angle between two (batches of) vectors. |
| Better precision than using the arccos dot product directly. |
| |
| ## Parameters |
| - `v1`: ndarray, shape (..., D): the first vector. |
| - `v2`: ndarray, shape (..., D): the second vector. |
| |
| ## Returns |
| `angle`: ndarray, shape (...): the angle between the two vectors. |
| """ |
| n1 = lite_norm(v1, axis=-1) |
| n2 = lite_norm(v2, axis=-1) |
| v1 = v1 / np.where(n1 == 0, 1, n1)[..., None] |
| v2 = v2 / np.where(n2 == 0, 1, n2)[..., None] |
| cos = lite_dot(v1, v2) |
| sin = np.minimum(lite_norm(v2 - v1 * cos[..., None], axis=-1), lite_norm(v1 - v2 * cos[..., None], axis=-1)) |
| return np.atan2(sin, cos) |
|
|
