Upload visualization/MEI138/geometry.py with huggingface_hub

visualization/MEI138/geometry.py

@@ -0,0 +1,303 @@
+"""Rotation conversion utilities for MEI representation.
+
+All functions support arbitrary batch dimensions (..., ).
+Coordinate convention: Y-up, right-handed (X-right, Y-up, Z-forward).
+
+Core rotation functions ported from HY-Motion (hymotion/utils/geometry.py),
+torch -> numpy.
+"""
+
+import numpy as np
+
+
+# ============================================================
+# Helpers (from HY-Motion)
+# ============================================================
+
+def _sqrt_positive_part(x: np.ndarray) -> np.ndarray:
+    """Returns np.sqrt(np.maximum(0, x))."""
+    ret = np.zeros_like(x)
+    positive_mask = x > 0
+    ret[positive_mask] = np.sqrt(x[positive_mask])
+    return ret
+
+
+def standardize_quaternion(quaternions: np.ndarray) -> np.ndarray:
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part first,
+            as array of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as array of shape (..., 4).
+    """
+    return np.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
+
+
+# ============================================================
+# Axis-angle <-> Quaternion <-> Rotation matrix
+# ============================================================
+
+def axis_angle_to_quaternion(axis_angle: np.ndarray) -> np.ndarray:
+    """Convert rotations given as axis/angle to quaternions.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as an array of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        quaternions with real part first, as array of shape (..., 4).
+    """
+    angles = np.linalg.norm(axis_angle, axis=-1, keepdims=True)
+    half_angles = angles * 0.5
+    # sin(angle/2) / angle, exact; limit -> 0.5 as angle -> 0
+    nonzero = angles != 0
+    safe_angles = np.where(nonzero, angles, np.ones_like(angles))
+    sin_half_angles_over_angles = np.where(
+        nonzero, np.sin(half_angles) / safe_angles, 0.5
+    )
+    quaternions = np.concatenate(
+        [np.cos(half_angles), axis_angle * sin_half_angles_over_angles], axis=-1
+    )
+    return quaternions
+
+
+def quaternion_to_matrix(quaternions: np.ndarray) -> np.ndarray:
+    """Convert rotations given as quaternions to rotation matrices.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as array of shape (..., 4).
+
+    Returns:
+        Rotation matrices as array of shape (..., 3, 3).
+    """
+    r, i, j, k = (
+        quaternions[..., 0],
+        quaternions[..., 1],
+        quaternions[..., 2],
+        quaternions[..., 3],
+    )
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = np.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        axis=-1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def axis_angle_to_matrix(axis_angle: np.ndarray) -> np.ndarray:
+    """Convert rotations given as axis/angle to rotation matrices.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as an array of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        Rotation matrices as array of shape (..., 3, 3).
+    """
+    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
+
+
+def matrix_to_quaternion(matrix: np.ndarray) -> np.ndarray:
+    """Convert rotations given as rotation matrices to quaternions.
+
+    Args:
+        matrix: Rotation matrices as array of shape (..., 3, 3).
+
+    Returns:
+        quaternions with real part first, as array of shape (..., 4).
+    """
+    if matrix.shape[-1] != 3 or matrix.shape[-2] != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+
+    batch_dim = matrix.shape[:-2]
+    m00, m01, m02, m10, m11, m12, m20, m21, m22 = np.split(
+        matrix.reshape(batch_dim + (9,)), 9, axis=-1
+    )
+    m00 = m00[..., 0]
+    m01 = m01[..., 0]
+    m02 = m02[..., 0]
+    m10 = m10[..., 0]
+    m11 = m11[..., 0]
+    m12 = m12[..., 0]
+    m20 = m20[..., 0]
+    m21 = m21[..., 0]
+    m22 = m22[..., 0]
+
+    q_abs = _sqrt_positive_part(
+        np.stack(
+            [
+                1.0 + m00 + m11 + m22,
+                1.0 + m00 - m11 - m22,
+                1.0 - m00 + m11 - m22,
+                1.0 - m00 - m11 + m22,
+            ],
+            axis=-1,
+        )
+    )
+
+    # we produce the desired quaternion multiplied by each of r, i, j, k
+    quat_by_rijk = np.stack(
+        [
+            np.stack(
+                [q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], axis=-1
+            ),
+            np.stack(
+                [m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], axis=-1
+            ),
+            np.stack(
+                [m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], axis=-1
+            ),
+            np.stack(
+                [m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], axis=-1
+            ),
+        ],
+        axis=-2,
+    )
+
+    # We floor here at 0.1 but the exact level is not important; if q_abs is small,
+    # the candidate won't be picked.
+    flr = 0.1
+    quat_candidates = quat_by_rijk / (2.0 * np.maximum(q_abs[..., None], flr))
+
+    # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
+    # forall i; we pick the best-conditioned one (with the largest denominator)
+    best = q_abs.argmax(axis=-1)  # (*batch_dim,)
+    # Advanced indexing to select the best candidate per element
+    flat_candidates = quat_candidates.reshape(-1, 4, 4)
+    flat_best = best.reshape(-1)
+    out = flat_candidates[np.arange(flat_candidates.shape[0]), flat_best, :]
+    out = out.reshape(batch_dim + (4,))
+    return standardize_quaternion(out)
+
+
+def quaternion_to_axis_angle(quaternions: np.ndarray) -> np.ndarray:
+    """Convert rotations given as quaternions to axis/angle.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as array of shape (..., 4).
+
+    Returns:
+        Rotations given as a vector in axis angle form, as an array
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    norms = np.linalg.norm(quaternions[..., 1:], axis=-1, keepdims=True)
+    half_angles = np.arctan2(norms, quaternions[..., :1])
+    angles = 2 * half_angles
+    # sin(half_angle) / angle, exact; limit -> 0.5 as angle -> 0
+    nonzero = angles != 0
+    safe_angles = np.where(nonzero, angles, np.ones_like(angles))
+    sin_half_angles_over_angles = np.where(
+        nonzero, np.sin(half_angles) / safe_angles, 0.5
+    )
+    return quaternions[..., 1:] / sin_half_angles_over_angles
+
+
+def matrix_to_axis_angle(matrix: np.ndarray) -> np.ndarray:
+    """Convert rotations given as rotation matrices to axis/angle.
+
+    Args:
+        matrix: Rotation matrices as array of shape (..., 3, 3).
+
+    Returns:
+        Rotations given as a vector in axis angle form, as an array
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
+
+
+# ============================================================
+# 6D continuous rotation representation (Zhou et al., CVPR 2019)
+# ============================================================
+
+def rotation_6d_to_matrix(rot6d: np.ndarray) -> np.ndarray:
+    """Convert 6D rotation representation to 3x3 rotation matrix.
+
+    Based on Zhou et al., "On the Continuity of Rotation Representations
+    in Neural Networks", CVPR 2019.
+
+    Args:
+        rot6d: array of shape (*, 6) of 6d rotation representations.
+
+    Returns:
+        rotation matrices of size (*, 3, 3).
+    """
+    x = rot6d.reshape(*rot6d.shape[:-1], 3, 2)
+    a1 = x[..., 0]
+    a2 = x[..., 1]
+    b1 = a1 / np.maximum(np.linalg.norm(a1, axis=-1, keepdims=True), 1e-12)
+    b2 = a2 - np.sum(b1 * a2, axis=-1, keepdims=True) * b1
+    b2 = b2 / np.maximum(np.linalg.norm(b2, axis=-1, keepdims=True), 1e-12)
+    b3 = np.cross(b1, b2, axis=-1)
+    return np.stack((b1, b2, b3), axis=-1)
+
+
+def matrix_to_rotation_6d(matrix: np.ndarray) -> np.ndarray:
+    """Convert 3x3 rotation matrix to 6D rotation representation.
+
+    Args:
+        matrix: rotation matrices of shape (*, 3, 3).
+
+    Returns:
+        6D rotation representation of shape (*, 6).
+    """
+    v1 = matrix[..., 0:1]
+    v2 = matrix[..., 1:2]
+    rot6d = np.concatenate([v1, v2], axis=-1).reshape(*matrix.shape[:-2], 6)
+    return rot6d
+
+
+# ============================================================
+# Yaw (Y-axis) rotation helpers (MEI-specific)
+# ============================================================
+
+def yaw_rotation_matrix(angle: np.ndarray) -> np.ndarray:
+    """Create rotation matrices for yaw (Y-axis rotation).
+
+    R_y(theta) maps local Z-forward to the heading direction in world XZ plane.
+
+    Args:
+        angle: (...) yaw angles in radians.
+
+    Returns:
+        R: (..., 3, 3) rotation matrices.
+    """
+    c = np.cos(angle)
+    s = np.sin(angle)
+    z = np.zeros_like(angle)
+    o = np.ones_like(angle)
+
+    return np.stack([
+        c,  z, s,
+        z,  o, z,
+        -s, z, c,
+    ], axis=-1).reshape(*angle.shape, 3, 3)
+
+
+def wrap_angle(angle: np.ndarray) -> np.ndarray:
+    """Wrap angle to [-pi, pi]."""
+    return (angle + np.pi) % (2 * np.pi) - np.pi