zed-playground/py_workspace/tests/test_alignment.py

import numpy as np
import pytest
from aruco.alignment import (
    compute_face_normal,
    rotation_align_vectors,
    apply_alignment_to_pose,
    get_face_normal_from_geometry,
    detect_ground_face,
)


def test_compute_face_normal_valid_quad():
    # Define a quad in the XY plane (normal should be Z)
    corners = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], dtype=np.float64)

    # v1 = corners[1] - corners[0] = [1, 0, 0]
    # v2 = corners[3] - corners[0] = [0, 1, 0]
    # normal = [1, 0, 0] x [0, 1, 0] = [0, 0, 1]

    normal = compute_face_normal(corners)
    np.testing.assert_allclose(normal, np.array([0, 0, 1]), atol=1e-10)


def test_compute_face_normal_valid_triangle():
    corners = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.float64)

    normal = compute_face_normal(corners)
    np.testing.assert_allclose(normal, np.array([0, 0, 1]), atol=1e-10)


def test_compute_face_normal_degenerate():
    # Collinear points
    corners = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0]], dtype=np.float64)
    with pytest.raises(ValueError, match="collinear or degenerate"):
        compute_face_normal(corners)

    # Too few points
    corners_few = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float64)
    with pytest.raises(ValueError, match="At least 3 corners"):
        compute_face_normal(corners_few)


def test_rotation_align_vectors_identity():
    v1 = np.array([0, 0, 1], dtype=np.float64)
    v2 = np.array([0, 0, 1], dtype=np.float64)
    R = rotation_align_vectors(v1, v2)
    np.testing.assert_allclose(R, np.eye(3), atol=1e-10)


def test_rotation_align_vectors_90_deg():
    v1 = np.array([1, 0, 0], dtype=np.float64)
    v2 = np.array([0, 1, 0], dtype=np.float64)
    R = rotation_align_vectors(v1, v2)

    # Check that R @ v1 == v2
    np.testing.assert_allclose(R @ v1, v2, atol=1e-10)
    # Check orthogonality
    np.testing.assert_allclose(R @ R.T, np.eye(3), atol=1e-10)


def test_rotation_align_vectors_antiparallel():
    v1 = np.array([0, 0, 1], dtype=np.float64)
    v2 = np.array([0, 0, -1], dtype=np.float64)
    R = rotation_align_vectors(v1, v2)

    np.testing.assert_allclose(R @ v1, v2, atol=1e-10)
    np.testing.assert_allclose(R @ R.T, np.eye(3), atol=1e-10)


def test_apply_alignment_to_pose():
    # Identity pose
    T = np.eye(4, dtype=np.float64)
    # 90 deg rotation around Z
    R_align = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=np.float64)

    T_aligned = apply_alignment_to_pose(T, R_align)

    expected = np.eye(4)
    expected[:3, :3] = R_align
    np.testing.assert_allclose(T_aligned, expected, atol=1e-10)

    # Non-identity pose
    T_pose = np.eye(4)
    T_pose[:3, 3] = [1, 2, 3]
    T_aligned_pose = apply_alignment_to_pose(T_pose, R_align)

    # Pre-multiplication: R_align @ T_pose
    # Rotation part: R_align @ I = R_align
    # Translation part: R_align @ [1, 2, 3] = [-2, 1, 3]
    expected_pose = np.eye(4)
    expected_pose[:3, :3] = R_align
    expected_pose[:3, 3] = R_align @ [1, 2, 3]

    np.testing.assert_allclose(T_aligned_pose, expected_pose, atol=1e-10)


def test_get_face_normal_from_geometry():
    face_marker_map = {"top": [1, 2]}
    # Marker 1: XY plane (normal Z)
    # Marker 2: XY plane (normal Z)
    marker_geometry = {
        1: np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], dtype=np.float64),
        2: np.array([[2, 2, 0], [3, 2, 0], [3, 3, 0], [2, 3, 0]], dtype=np.float64),
    }

    normal = get_face_normal_from_geometry("top", marker_geometry, face_marker_map)
    assert normal is not None
    np.testing.assert_allclose(normal, np.array([0, 0, 1]), atol=1e-10)

    # Missing face map
    assert get_face_normal_from_geometry("top", marker_geometry, None) is None

    # Missing face name
    assert (
        get_face_normal_from_geometry("bottom", marker_geometry, face_marker_map)
        is None
    )

    # No visible markers for face
    assert (
        get_face_normal_from_geometry("top", {3: marker_geometry[1]}, face_marker_map)
        is None
    )


def test_detect_ground_face():
    face_marker_map = {
        "bottom": [1],
        "top": [2],
    }
    # Marker 1: normal [0, -1, 0] (aligned with camera up [0, -1, 0])
    # Marker 2: normal [0, 1, 0] (opposite to camera up)
    marker_geometry = {
        1: np.array(
            [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]], dtype=np.float64
        ),  # Normal is [0, 1, 0] or [0, -1, 0]?
        # v1 = [1, 0, 0], v2 = [0, 0, 1] -> cross = [0, -1, 0]
        2: np.array([[0, 1, 0], [1, 1, 0], [1, 1, -1], [0, 1, -1]], dtype=np.float64),
        # v1 = [1, 0, 0], v2 = [0, 0, -1] -> cross = [0, 1, 0]
    }

    camera_up = np.array([0, -1, 0], dtype=np.float64)

    # Both visible
    res = detect_ground_face({1, 2}, marker_geometry, camera_up, face_marker_map)
    assert res is not None
    face_name, normal = res
    assert face_name == "bottom"
    np.testing.assert_allclose(normal, np.array([0, -1, 0]), atol=1e-10)

    # Only top visible
    res = detect_ground_face({2}, marker_geometry, camera_up, face_marker_map)
    assert res is not None
    face_name, normal = res
    assert face_name == "top"
    np.testing.assert_allclose(normal, np.array([0, 1, 0]), atol=1e-10)

    # None visible
    assert (
        detect_ground_face(set(), marker_geometry, camera_up, face_marker_map) is None
    )

    # Missing map
    assert detect_ground_face({1, 2}, marker_geometry, camera_up, None) is None