vid2rig / tests /test_rotutil.py
BladeSzaSza's picture
Upload folder using huggingface_hub
3d348c4 verified
Raw
History Blame Contribute Delete
933 Bytes
import numpy as np
from scipy.spatial.transform import Rotation as R
from vid2rig.rotutil import aa_to_quat, fix_hemisphere, slerp_series
def test_aa_to_quat_matches_scipy():
aa = np.array([[0, 0, np.pi / 2], [0.3, -0.2, 0.1]])
assert np.allclose(np.abs(aa_to_quat(aa)), np.abs(R.from_rotvec(aa).as_quat()))
def test_fix_hemisphere_flips_sign():
q = R.from_rotvec([0, 0, 0.1]).as_quat()
series = np.stack([q, -q, q]) # artificially flipped middle frame
fixed = fix_hemisphere(series)
dots = np.sum(fixed[1:] * fixed[:-1], axis=1)
assert np.all(dots >= -1e-6)
def test_slerp_midpoint():
q0 = R.from_rotvec([0, 0, 0.0]).as_quat()
q1 = R.from_rotvec([0, 0, np.pi / 2]).as_quat()
out = slerp_series(np.stack([q0, q1]), np.array([0.0, 1.0]), np.array([0.5]))
expected = R.from_rotvec([0, 0, np.pi / 4]).as_quat()
assert np.allclose(np.abs(out[0]), np.abs(expected), atol=1e-6)