mirror of https://github.com/commaai/openpilot.git
53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
import numpy as np
|
|
from collections.abc import Callable
|
|
|
|
from openpilot.common.transformations.transformations import (ecef_euler_from_ned_single,
|
|
euler2quat_single,
|
|
euler2rot_single,
|
|
ned_euler_from_ecef_single,
|
|
quat2euler_single,
|
|
quat2rot_single,
|
|
rot2euler_single,
|
|
rot2quat_single)
|
|
|
|
|
|
def numpy_wrap(function, input_shape, output_shape) -> Callable[..., np.ndarray]:
|
|
"""Wrap a function to take either an input or list of inputs and return the correct shape"""
|
|
def f(*inps):
|
|
*args, inp = inps
|
|
inp = np.array(inp)
|
|
shape = inp.shape
|
|
|
|
if len(shape) == len(input_shape):
|
|
out_shape = output_shape
|
|
else:
|
|
out_shape = (shape[0],) + output_shape
|
|
|
|
# Add empty dimension if inputs is not a list
|
|
if len(shape) == len(input_shape):
|
|
inp.shape = (1, ) + inp.shape
|
|
|
|
result = np.asarray([function(*args, i) for i in inp])
|
|
result.shape = out_shape
|
|
return result
|
|
return f
|
|
|
|
|
|
euler2quat = numpy_wrap(euler2quat_single, (3,), (4,))
|
|
quat2euler = numpy_wrap(quat2euler_single, (4,), (3,))
|
|
quat2rot = numpy_wrap(quat2rot_single, (4,), (3, 3))
|
|
rot2quat = numpy_wrap(rot2quat_single, (3, 3), (4,))
|
|
euler2rot = numpy_wrap(euler2rot_single, (3,), (3, 3))
|
|
rot2euler = numpy_wrap(rot2euler_single, (3, 3), (3,))
|
|
ecef_euler_from_ned = numpy_wrap(ecef_euler_from_ned_single, (3,), (3,))
|
|
ned_euler_from_ecef = numpy_wrap(ned_euler_from_ecef_single, (3,), (3,))
|
|
|
|
quats_from_rotations = rot2quat
|
|
quat_from_rot = rot2quat
|
|
rotations_from_quats = quat2rot
|
|
rot_from_quat = quat2rot
|
|
euler_from_rot = rot2euler
|
|
euler_from_quat = quat2euler
|
|
rot_from_euler = euler2rot
|
|
quat_from_euler = euler2quat
|