ylff/utils/coordinate_utils.py

"""
Coordinate system conversion utilities.
ARKit uses Y-up, right-handed.
DA3/COLMAP typically use Z-up, right-handed (OpenCV convention).
"""

import numpy as np


def arkit_to_opencv_transform() -> np.ndarray:
    """
    Transform from ARKit coordinate system to OpenCV/COLMAP convention.

    ARKit: Y-up, right-handed
    OpenCV: Z-up, right-handed

    Conversion: Rotate -90° around X axis (Y -> Z)
    """
    # Rotation matrix: -90° around X axis
    # This maps: Y -> Z, Z -> -Y
    R = np.array(
        [
            [1, 0, 0],
            [0, 0, 1],
            [0, -1, 0],
        ]
    )
    return R


def convert_arkit_to_opencv(pose_c2w: np.ndarray) -> np.ndarray:
    """
    Convert ARKit camera-to-world pose to OpenCV/COLMAP convention.

    Args:
        pose_c2w: (4, 4) ARKit c2w pose (Y-up)

    Returns:
        (4, 4) OpenCV c2w pose (Z-up)
    """
    if pose_c2w.shape != (4, 4):
        raise ValueError(f"Expected (4, 4) pose, got {pose_c2w.shape}")

    R_convert = arkit_to_opencv_transform()

    # Extract rotation and translation
    R_arkit = pose_c2w[:3, :3]
    t_arkit = pose_c2w[:3, 3]

    # Convert rotation: R_opencv = R_convert @ R_arkit @ R_convert^T
    R_opencv = R_convert @ R_arkit @ R_convert.T

    # Convert translation: t_opencv = R_convert @ t_arkit
    t_opencv = R_convert @ t_arkit

    # Build new pose
    pose_opencv = np.eye(4)
    pose_opencv[:3, :3] = R_opencv
    pose_opencv[:3, 3] = t_opencv

    return pose_opencv


def convert_opencv_to_arkit(pose_c2w: np.ndarray) -> np.ndarray:
    """
    Convert OpenCV/COLMAP camera-to-world pose to ARKit convention.

    Args:
        pose_c2w: (4, 4) OpenCV c2w pose (Z-up)

    Returns:
        (4, 4) ARKit c2w pose (Y-up)
    """
    R_convert = arkit_to_opencv_transform()
    R_convert_inv = R_convert.T  # Inverse rotation

    R_opencv = pose_c2w[:3, :3]
    t_opencv = pose_c2w[:3, 3]

    # Convert rotation: R_arkit = R_convert_inv @ R_opencv @ R_convert
    R_arkit = R_convert_inv @ R_opencv @ R_convert

    # Convert translation: t_arkit = R_convert_inv @ t_opencv
    t_arkit = R_convert_inv @ t_opencv

    pose_arkit = np.eye(4)
    pose_arkit[:3, :3] = R_arkit
    pose_arkit[:3, 3] = t_arkit

    return pose_arkit


def convert_arkit_c2w_to_w2c(pose_c2w: np.ndarray, convert_coords: bool = True) -> np.ndarray:
    """
    Convert ARKit c2w pose to w2c (for DA3 compatibility).

    Args:
        pose_c2w: (4, 4) ARKit c2w pose
        convert_coords: If True, convert from ARKit to OpenCV coordinate system

    Returns:
        (3, 4) w2c pose in OpenCV convention (if convert_coords=True) or ARKit convention
    """
    if convert_coords:
        pose_c2w = convert_arkit_to_opencv(pose_c2w)

    # Invert to get w2c
    pose_w2c = np.linalg.inv(pose_c2w)

    # Extract 3x4
    return pose_w2c[:3, :]


def test_coordinate_conversion():
    """Test coordinate conversion."""
    # Create a test ARKit pose (Y-up)
    # Camera at origin, looking down +Z (in ARKit convention)
    pose_arkit = np.eye(4)
    pose_arkit[:3, 3] = [1.0, 2.0, 3.0]  # Translation

    print("ARKit pose (Y-up):")
    print(pose_arkit)
    print(f"  Translation: {pose_arkit[:3, 3]}")

    # Convert to OpenCV
    pose_opencv = convert_arkit_to_opencv(pose_arkit)
    print("\nOpenCV pose (Z-up):")
    print(pose_opencv)
    print(f"  Translation: {pose_opencv[:3, 3]}")

    # Convert back
    pose_arkit_back = convert_opencv_to_arkit(pose_opencv)
    print("\nARKit pose (converted back):")
    print(pose_arkit_back)
    print(f"  Translation: {pose_arkit_back[:3, 3]}")

    # Check round-trip
    assert np.allclose(pose_arkit, pose_arkit_back), "Round-trip conversion failed!"
    print("\n✓ Round-trip conversion successful!")


if __name__ == "__main__":
    test_coordinate_conversion()