add rotation matrix -> quaternion and yaw -> quaternion conversion ut…

…ility functions

argoverse/utils/transform.py

@@ -16,6 +16,26 @@
 logger = logging.getLogger(__name__)
 
 
+def yaw_to_quaternion3d(yaw: float) -> np.ndarray:
+    """Convert a rotation angle in the xy plane (i.e. about the z axis) to a quaternion.
+
+    Args:
+        yaw: angle to rotate about the z-axis, representing an Euler angle, in radians
+
+    Returns:
+        array w/ quaternion coefficients (qw,qx,qy,qz) in scalar-first order, per Argoverse convention.
+    """
+    qx, qy, qz, qw = Rotation.from_euler(seq="z", angles=yaw, degrees=False).as_quat()
+    return np.array([qw, qx, qy, qz])
+
+
+def rotmat2quat(R: np.ndarray) -> np.ndarray:
+    """Convert a rotation-matrix to a quaternion in Argo's scalar-first notation (w, x, y, z)."""
+    quat_xyzw = Rotation.from_matrix(R).as_quat()
+    quat_wxyz = quat_scipy2argo(quat_xyzw)
+    return quat_wxyz
+
+
 def quat2rotmat(q: np.ndarray) -> np.ndarray:
     """Normalizes a quaternion to unit-length, then converts it into a rotation matrix.
 
@@ -48,6 +68,13 @@ def quat_argo2scipy(q: np.ndarray) -> np.ndarray:
     return q_scipy
 
 
+def quat_scipy2argo(q: np.ndarray) -> np.ndarray:
+    """Re-order Scipy's scalar-last [x,y,z,w] quaternion order to Argoverse's scalar-first [w,x,y,z]."""
+    x, y, z, w = q
+    q_argo = np.array([w, x, y, z])
+    return q_argo
+
+
 def quat_argo2scipy_vectorized(q: np.ndarray) -> np.ndarray:
     """Re-order Argoverse's scalar-first [w,x,y,z] quaternion order to Scipy's scalar-last [x,y,z,w]"""
     return q[..., [1, 2, 3, 0]]