def rf_tree(self, register_rf_tree):
""""""
register_rf_tree(
tc1=(1.0, 0.0, 0.0),
tc2=(-1.0, 0.0, 0.0),
rc2=from_euler_angles(yaw=np.pi),
)
def test_qinv(self):
""""""
q = from_euler_angles(roll=np.pi / 4, return_quaternion=True)
assert qinv(q) == 1 / q
npt.assert_equal(qinv(as_float_array(q)), as_float_array(1 / q))
q_arr = np.tile(as_float_array(q), (10, 1)).T
npt.assert_equal(qinv(q_arr, 0)[:, 0], as_float_array(1 / q))
def test_transform_quaternions(self, rf_tree):
""""""
arr_child2 = (1.0, 0.0, 0.0, 0.0)
arr_exp = from_euler_angles(yaw=np.pi)
# tuple
arr_child1 = rbm.transform_quaternions(arr_child2,
into="child1",
outof="child2")
npt.assert_almost_equal(arr_child1, arr_exp)
def test_best_fit_rotation(self, get_rf_tree):
""""""
rf_world, rf_child1, _ = get_rf_tree(
(0.0, 0.0, 0.0), from_euler_angles(roll=np.pi / 2)
)
v1 = np.random.randn(10, 3)
v2 = rf_world.transform_points(v1, rf_child1)
r = best_fit_rotation(v1, v2)
_, r_exp, _ = rf_world.lookup_transform(rf_child1)
npt.assert_allclose(np.abs(r), np.abs(r_exp), rtol=1.0, atol=1e-10)
def test_qmean(self):
""""""
q = np.hstack((
from_euler_angles(roll=np.pi / 4, return_quaternion=True),
from_euler_angles(roll=-np.pi / 4, return_quaternion=True),
from_euler_angles(pitch=np.pi / 4, return_quaternion=True),
from_euler_angles(pitch=-np.pi / 4, return_quaternion=True),
quaternion(1.0, 0.0, 0.0, 0.0),
))
# quaternion dtype
qm = qmean(q)
npt.assert_allclose(as_float_array(qm), np.array([1.0, 0.0, 0.0, 0.0]))
# float dtype
qm = qmean(as_float_array(q).T, qaxis=0)
npt.assert_allclose(qm, np.array([1.0, 0.0, 0.0, 0.0]))
# not all axes
qm = qmean(np.tile(q, (10, 1)), axis=1)
npt.assert_allclose(
as_float_array(qm),
np.tile(np.array([1.0, 0.0, 0.0, 0.0]), (10, 1)),
)
def test_qmul(self):
""""""
q1 = from_euler_angles(roll=np.pi / 4, return_quaternion=True)
q2 = from_euler_angles(roll=np.pi / 4, return_quaternion=True)
assert qmul(q1, q1) == q1 * q2
npt.assert_equal(
qmul(as_float_array(q1), as_float_array(q2)),
as_float_array(q1 * q2),
)
q_arr = np.tile(as_float_array(q1), (10, 1))
npt.assert_equal(
qmul(q_arr, as_float_array(q2)),
as_float_array(as_quat_array(q_arr) * q2),
)
# only one quaternion
with pytest.raises(ValueError):
qmul(q1)
# different dtypes
with pytest.raises(ValueError):
qmul(q1, as_float_array(q2))
def test_best_fit_transform_xr(self, get_rf_tree):
""""""
xr = pytest.importorskip("xarray")
rf_world, rf_child1, _ = get_rf_tree(
(1.0, 0.0, 0.0), from_euler_angles(roll=np.pi / 2)
)
v1 = np.random.randn(10, 3)
v2 = rf_world.transform_points(v1, rf_child1)
v1_da = xr.DataArray(
v1, {"cartesian_axis": ["x", "y", "z"]}, ("time", "cartesian_axis")
)
t, r = best_fit_transform(v1_da.T, v2.T, dim="cartesian_axis")
assert t.dims == ("cartesian_axis",)
assert r.dims == ("quaternion_axis",)
def test_iterative_closest_point(self, get_rf_tree):
""""""
np.random.seed(42)
rf_world, rf_child1, _ = get_rf_tree(
(1.0, 0.0, 0.0), from_euler_angles(yaw=np.pi / 6)
)
x, y = np.meshgrid(np.arange(10), np.arange(20))
v1 = np.column_stack(
(x.flatten() ** 2, y.flatten() ** 2, np.ones(x.size))
) + 0.01 * np.random.randn(x.size, 3)
v2 = np.random.permutation(
rf_world.transform_points(v1, rf_child1)[1:]
)
t, r = iterative_closest_point(v1, v2)
t_exp, r_exp, _ = rf_world.lookup_transform(rf_child1)
npt.assert_allclose(t, t_exp, rtol=1.0, atol=0.01)
npt.assert_allclose(np.abs(r), np.abs(r_exp), rtol=1.0, atol=1e-4)
def test_transform_quaternions_xr(self, rf_tree):
""""""
xr = pytest.importorskip("xarray")
arr_child2 = (1.0, 0.0, 0.0, 0.0)
arr_exp = from_euler_angles(yaw=np.pi)
da_child2 = xr.DataArray(
np.tile(arr_child2, (10, 1)),
{"time": np.arange(10)},
("time", "quaternion_axis"),
attrs={"reference_frame": "child2"},
)
da_child1 = rbm.transform_quaternions(
da_child2,
into="child1",
dim="quaternion_axis",
timestamps="time",
)
assert da_child1.shape == (10, 4)
assert da_child1.attrs["reference_frame"] == "child1"
assert da_child1.attrs["representation_frame"] == "child1"
npt.assert_almost_equal(da_child1[0], arr_exp)
# multi-dimensional vectors
da_child2 = xr.DataArray(
np.tile(arr_child2, (5, 10, 1)),
{"time": np.arange(10)},
("extra_dim", "time", "quaternion_axis"),
)
da_child1 = rbm.transform_quaternions(
da_child2,
into="child1",
outof="child2",
dim="quaternion_axis",
timestamps="time",
)
assert da_child1.shape == (5, 10, 4)
assert da_child1.dims == ("extra_dim", "time", "quaternion_axis")
npt.assert_almost_equal(da_child1[0, 0], arr_exp)
def test_rotate_vectors(self):
""""""
v = np.ones((10, 3))
q = np.tile(from_euler_angles(yaw=np.pi / 4, return_quaternion=True),
10)
vr = np.vstack((np.zeros(10), np.sqrt(2) * np.ones(10), np.ones(10))).T
# single quaternion, single vector
vr_act = rotate_vectors(q[0], v[0])
np.testing.assert_allclose(vr[0], vr_act, rtol=1.0)
# single quaternion, multiple vectors
vr_act = rotate_vectors(q[0], v)
np.testing.assert_allclose(vr, vr_act, rtol=1.0)
# single quaternion, explicit axis
vr_act = rotate_vectors(q[0], v, axis=1)
np.testing.assert_allclose(vr, vr_act, rtol=1.0)
# multiple quaternions, multiple vectors
vr_act = rotate_vectors(q, v)
np.testing.assert_allclose(vr, vr_act)
# different axis
vr_act = rotate_vectors(q, v.T, axis=0)
np.testing.assert_allclose(vr.T, vr_act)
# singleton expansion
vr_act = rotate_vectors(q[:, None], v[None, ...])
np.testing.assert_allclose(np.tile(vr, (10, 1, 1)), vr_act)
# float dtype
vr_act = rotate_vectors(as_float_array(q[0]), v[0])
np.testing.assert_allclose(vr[0], vr_act, rtol=1.0)
with pytest.raises(ValueError):
rotate_vectors(q, v.T)
with pytest.raises(ValueError):
rotate_vectors(q, np.ones((10, 4)))
def test_best_fit_transform(self, get_rf_tree, icp_test_data):
""""""
rf_world, rf_child1, _ = get_rf_tree(
(1.0, 0.0, 0.0), from_euler_angles(roll=np.pi / 2)
)
v1 = np.random.randn(10, 3)
v2 = rf_world.transform_points(v1, rf_child1)
t, r = best_fit_transform(v1, v2)
t_exp, r_exp, _ = rf_world.lookup_transform(rf_child1)
npt.assert_allclose(t, t_exp, rtol=1.0, atol=1e-10)
npt.assert_allclose(np.abs(r), np.abs(r_exp), rtol=1.0, atol=1e-10)
# real data
v1 = icp_test_data["t265"]
v2 = icp_test_data["vicon"]
t, r = best_fit_transform(v1, v2)
npt.assert_allclose(t, [-1.89872037, 0.61755277, 0.95930489])
npt.assert_allclose(
r,
[-6.87968663e-01, 4.73801246e-04, 9.12595868e-03, 7.25682859e-01],
)
def load_optitrack(filepath, import_format="numpy"):
""" Load rigid body poses from OptiTrack csv export.
Parameters
----------
filepath: str or path-like
Path to csv file.
import_format: {"numpy", "pandas", "xarray"}, default "numpy"
Import format for rigid body poses. "numpy" returns a (position,
orientation, timestamps) tuple for each body, "pandas" returns a
DataFrame and "xarray" returns a Dataset.
Returns
-------
data_dict: dict
Dictionary with one entry for each rigid body. See ``import_format``
for the format of each entry.
"""
try:
import pandas as pd
except ImportError:
raise ModuleNotFoundError("Install pandas to use optitrack import")
# parse header with metadata
header = pd.read_csv(filepath, nrows=1, header=None).values
meta = {k: v for k, v in zip(header[0][0::2], header[0][1::2])}
# parse body with poses of rigid bodies and marker positions
df = pd.read_csv(filepath,
header=[2, 3, 5, 6],
index_col=0,
skip_blank_lines=False)
# compute timestamps
timestamps = pd.to_datetime(
meta["Capture Start Time"],
format="%Y-%m-%d %I.%M.%S.%f %p") + pd.to_timedelta(df.values[:, 0],
unit="s")
# only import rigid bodies, not markers
df = df["Rigid Body"]
# set index to timestamps and multi-index levels
df = df.set_index(timestamps).rename_axis("time")
df.columns = df.columns.set_names(["body ID", "motion type", "axis"])
# split dataframe into dict with one entry for each rigid body
data_dict = {
idx: gp.xs(idx, level=0, axis=1)
for idx, gp in df.groupby(level=0, axis=1)
}
# convert cm to m (if applicable)
if meta.get("Length Units", None) == "Centimeters":
for name in data_dict:
data_dict[name]["Position"] /= 100
# convert orientation to quaternions
for name in data_dict:
rpy = np.deg2rad(data_dict[name]["Rotation"].values)
order = (meta.get("Rotation Type",
"XYZ").replace("X",
"r").replace("Y",
"p").replace("Z", "y"))
q = from_euler_angles(rpy, order=order)
position = data_dict[name]["Position"]
position.columns = position.columns.str.lower()
orientation = pd.DataFrame(q,
index=position.index,
columns=["w", "x", "y", "z"])
data_dict[name] = pd.concat(
{
"position": position,
"orientation": orientation
},
axis=1,
)
# return in requested format
if import_format == "pandas":
return data_dict
elif import_format == "numpy":
return {
k: (v["position"].values, v["orientation"].values, v.index.values)
for k, v in data_dict.items()
}
elif import_format == "xarray":
import xarray as xr
for name in data_dict:
position = xr.DataArray(
data_dict[name]["position"],
name="position",
dims=("time", "cartesian_axis"),
)
orientation = xr.DataArray(
data_dict[name].orientation,
name="orientation",
dims=("time", "quaternion_axis"),
)
data_dict[name] = xr.merge((position, orientation))
return data_dict
else:
raise ValueError(f"Unsupported import format: {import_format}")
def test_from_euler_angles(self):
""""""
rpy_r = (np.pi / 2, 0, 0)
q_r = (np.sqrt(2) / 2, np.sqrt(2) / 2, 0, 0)
q_rp = (0.5, 0.5, 0.5, 0.5)
q_ypr = (np.sqrt(2) / 2, 0, np.sqrt(2) / 2, 0)
# single rpy
q = from_euler_angles(rpy_r)
np.testing.assert_almost_equal(q, q_r)
# rpy array
q = from_euler_angles(np.tile(rpy_r, (10, 1)))
np.testing.assert_almost_equal(q, np.tile(q_r, (10, 1)))
# different axis
q = from_euler_angles(np.tile(rpy_r, (10, 1)).T, axis=0)
np.testing.assert_almost_equal(q, np.tile(q_r, (10, 1)).T)
# single roll
q = from_euler_angles(roll=np.pi / 2)
np.testing.assert_almost_equal(q, q_r)
# single roll, pitch & yaw
q = from_euler_angles(roll=np.pi / 2, pitch=np.pi / 2)
np.testing.assert_almost_equal(q, q_rp)
# single roll, pitch, yaw + different order
q = from_euler_angles(roll=np.pi / 2,
pitch=np.pi / 2,
yaw=np.pi / 2,
order="ypr")
np.testing.assert_almost_equal(q, q_ypr)
# no input
with pytest.raises(ValueError):
from_euler_angles()
# rpy and roll together
with pytest.raises(ValueError):
from_euler_angles(rpy_r, roll=np.pi / 2)
# inconsistent shapes
with pytest.raises(ValueError):
from_euler_angles(roll=np.ones(10), pitch=1)