add and test 2d2d solver

src/ouroboros_opengv/__init__.py

@@ -1 +1,77 @@
-from _ouroboros_opengv import *
+from typing import Optional
+
+import numpy as np
+from _ouroboros_opengv import Solver2d2d, Solver2d3d, solve_2d2d, solve_2d3d, solve_3d3d
+
+# TODO(nathan) use actual type alias once we move beyond 3.8
+# Matrix3d = np.ndarray[np.float64[3, 3]]
+Matrix3d = np.ndarray
+
+
+def inverse_camera_matrix(K: Matrix3d) -> Matrix3d:
+    """
+    Get inverse camera matrix.
+
+    Args:
+        K: Original camera matrix.
+
+    Returns:
+        Inverted camera matrix that takes pixel space coordinates to unit coordinates.
+    """
+    K_inv = np.eye(3)
+    K_inv[0, 0] = 1.0 / K[0, 0]
+    K_inv[1, 1] = 1.0 / K[1, 1]
+    K_inv[0, 2] = -K[0, 2] / K[0, 0]
+    K_inv[1, 2] = -K[1, 2] / K[1, 1]
+    return K_inv
+
+
+def get_bearings(K: Matrix3d, features: np.ndarray) -> np.ndarray:
+    """
+    Get bearings for undistorted features in pixel space.
+
+    Args:
+        K: Camera matrix for features.
+        features: Pixel coordinates in a 2xN matrix.
+
+    Returns:
+        Bearing vectors in a 3xN matrix.
+    """
+    K_inv = inverse_camera_matrix(K)
+    bearings = np.vstack((features, np.ones(features.shape[1])))
+    bearings = K_inv @ bearings
+    bearings /= np.linalg.norm(bearings, axis=0)
+    return bearings
+
+
+def recover_pose_opengv(
+    K_query: Matrix3d,
+    query_features: np.ndarray,
+    K_match: Matrix3d,
+    match_features: np.ndarray,
+    correspondences: np.ndarray,
+    solver=Solver2d2d.STEWENIUS,
+) -> Optional[np.ndarray]:
+    """
+    Recover pose up to scale from 2d correspondences.
+
+    Args:
+        K_query: Camera matrix for query features.
+        query_features: 2xN matrix of pixel features for query frame.
+        K_match: Camera matrix for match features.
+        match_features: 2xN matrix of pixel feature for match frame.
+        correspondences: Nx2 indices of feature matches (query -> match)
+        solver: Underlying 2d2d algorithm.
+
+    Returns:
+        match_T_query if underlying solver is successful.
+    """
+    query_bearings = get_bearings(K_query, query_features[:, correspondences[:, 0]])
+    match_bearings = get_bearings(K_match, match_features[:, correspondences[:, 1]])
+    # order is src (query), dest (match) for dest_T_src (match_T_query)
+    result = solve_2d2d(query_bearings, match_bearings, solver=solver)
+    if not result:
+        return None
+
+    match_T_query = result.dest_T_src
+    return match_T_query

tests/test_opengv_interface.py

@@ -0,0 +1,72 @@
+"""Test opengv solver."""
+
+import numpy as np
+import pytest
+
+import ouroboros_opengv as ogv
+
+
+def test_inverse_camera_matrix():
+    """Check that explicit inverse is correct."""
+    orig = np.array([[500.0, 0.0, 320.0], [0.0, 500.0, 240.0], [0.0, 0.0, 1.0]])
+    result = ogv.inverse_camera_matrix(orig)
+    assert result == pytest.approx(np.linalg.inv(orig))
+
+
+def test_bearings():
+    """Check that bearing math is correct."""
+    K = np.array([[10.0, 0.0, 5.0], [0.0, 5.0, 2.5], [0.0, 0.0, 1.0]])
+    features = np.array([[5.0, 2.5], [15.0, 2.5], [5.0, -2.5]]).T
+    bearings = ogv.get_bearings(K, features)
+    expected = np.array(
+        [
+            [0.0, 0.0, 1.0],
+            [1.0 / np.sqrt(2), 0.0, 1.0 / np.sqrt(2)],
+            [0.0, -1.0 / np.sqrt(2), 1.0 / np.sqrt(2)],
+        ]
+    ).T
+    assert bearings == pytest.approx(expected)
+
+
+def _shuffle_features(features):
+    indices = np.arange(features.shape[1])
+    np.random.shuffle(indices)
+    return indices, features[:, indices].copy()
+
+
+def test_solver():
+    """Test that two-view geometry is called correct."""
+    query_features = np.random.normal(size=(2, 100))
+    query_bearings = ogv.get_bearings(np.eye(3), query_features)
+
+    yaw = np.pi / 4.0
+    match_R_query = np.array(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, np.cos(yaw), -np.sin(yaw)],
+            [0.0, np.sin(yaw), np.cos(yaw)],
+        ]
+    )
+    match_t_query = np.array([1.0, -1.2, 0.8]).reshape((3, 1))
+    match_bearings = match_R_query @ query_bearings + match_t_query
+    match_features = match_bearings[:2, :] / match_bearings[2, :]
+
+    indices = np.arange(query_bearings.shape[1])
+    new_indices, match_features = _shuffle_features(match_features)
+
+    # needs to be query -> match (so need indices that were used by shuffle for query)
+    correspondences = np.vstack((new_indices, indices)).T
+
+    match_T_query = ogv.recover_pose_opengv(
+        np.eye(3),
+        query_features,
+        np.eye(3),
+        match_features,
+        correspondences,
+        solver=ogv.Solver2d2d.NISTER,
+    )
+
+    t_expected = np.squeeze(match_t_query / np.linalg.norm(match_t_query))
+    t_result = np.squeeze(match_T_query[:3, 3] / np.linalg.norm(match_T_query[:3, 3]))
+    assert match_T_query[:3, :3] == pytest.approx(match_R_query)
+    assert t_result == pytest.approx(t_expected)