enable zero-copy on opengv side

MIT-SPARK · Feb 6, 2025 · 9b7b950 · 9b7b950
1 parent 9dd0c9a
commit 9b7b950
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 120 deletions.
diff --git a/.clang-format b/.clang-format
@@ -4,7 +4,7 @@ BasedOnStyle: Google
 Language: Cpp
 DerivePointerAlignment: false
 PointerAlignment: Left
-ColumnLimit: 88
+ColumnLimit: 120
 BinPackArguments: false
 BinPackParameters: false
 ---
diff --git a/src/ouroboros_opengv/__init__.py b/src/ouroboros_opengv/__init__.py
@@ -1,6 +1,7 @@
 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
@@ -32,15 +33,16 @@ def get_bearings(K: Matrix3d, features: np.ndarray) -> np.ndarray:
 
     Args:
         K: Camera matrix for features.
-        features: Pixel coordinates in a 2xN matrix.
+        features: Pixel coordinates in a Nx2 matrix.
 
     Returns:
-        Bearing vectors in a 3xN matrix.
+        Bearing vectors in a Nx3 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)
+    bearings = np.hstack((features, np.ones((features.shape[0], 1))))
+    bearings = bearings @ K_inv.T
+    # for broadcasting to be correct needs to be [N, 1] to divide rowwise
+    bearings /= np.linalg.norm(bearings, axis=1)[..., np.newaxis]
     return bearings
 
 
@@ -66,10 +68,10 @@ def recover_pose_opengv(
     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]])
+    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)
+    result = solve_2d2d(query_bearings.T, match_bearings.T, solver=solver)
     if not result:
         return None
 

diff --git a/src/ouroboros_opengv/_bindings.cpp b/src/ouroboros_opengv/_bindings.cpp
@@ -16,11 +16,11 @@ using opengv::relative_pose::RelativeAdapterBase;
 using opengv::sac::Ransac;
 using opengv::sac_problems::point_cloud::PointCloudSacProblem;
 
-using RelativePoseProblem =
-    opengv::sac_problems::relative_pose::CentralRelativePoseSacProblem;
+using RelativePoseProblem = opengv::sac_problems::relative_pose::CentralRelativePoseSacProblem;
 using AbsolutePoseProblem = opengv::sac_problems::absolute_pose::AbsolutePoseSacProblem;
 
 using Vec3Matrix = Eigen::Matrix<double, 3, Eigen::Dynamic>;
+using PyVec3Matrix = Eigen::Ref<const Vec3Matrix>;
 
 inline void checkInputs(const Vec3Matrix& m1,
                         const Vec3Matrix& m2,
@@ -31,100 +31,65 @@ inline void checkInputs(const Vec3Matrix& m1,
   }
 
   std::stringstream ss;
-  ss << "Invalid input shapes! Given " << m1_name << ": [" << m1.rows() << ", "
-     << m1.cols() << "] and " << m2_name << ": [" << m2.rows() << ", " << m2.cols()
-     << "] (columns do not match: " << m1.cols() << " != " << m2.cols() << ")";
+  ss << "Invalid input shapes! Given " << m1_name << ": [" << m1.rows() << ", " << m1.cols() << "] and " << m2_name
+     << ": [" << m2.rows() << ", " << m2.cols() << "] (columns do not match: " << m1.cols() << " != " << m2.cols()
+     << ")";
   throw std::invalid_argument(ss.str());
 }
 
+// NOTE(nathan) we don't both exposing priors for 2d-2d
 struct EigenRelativeAdaptor : public RelativeAdapterBase {
-  EigenRelativeAdaptor(const Vec3Matrix& src, const Vec3Matrix& dest)
-      : EigenRelativeAdaptor(src, dest, Eigen::Matrix3d::Identity()) {}
-
-  EigenRelativeAdaptor(const Vec3Matrix& src,
-                       const Vec3Matrix& dest,
-                       const Eigen::Matrix3d& rotation_prior)
-      : EigenRelativeAdaptor(src, dest, rotation_prior, Eigen::Vector3d::Zero()) {}
-
-  EigenRelativeAdaptor(const Vec3Matrix& _src,
-                       const Vec3Matrix& _dest,
-                       const Eigen::Matrix3d& rotation_prior,
-                       const Eigen::Vector3d& translation_prior)
-      : RelativeAdapterBase(translation_prior, rotation_prior), src(_src), dest(_dest) {
+  EigenRelativeAdaptor(const PyVec3Matrix _src, const PyVec3Matrix _dest)
+      : RelativeAdapterBase(Eigen::Vector3d::Zero(), Eigen::Matrix3d::Identity()), src(_src), dest(_dest) {
     checkInputs(src, dest, "src_bearings", "dest_bearings");
   }
 
-  Eigen::Vector3d getBearingVector1(size_t index) const override {
-    return dest.block<3, 1>(0, index);
-  }
+  Eigen::Vector3d getBearingVector1(size_t index) const override { return dest.block<3, 1>(0, index); }
 
-  Eigen::Vector3d getBearingVector2(size_t index) const override {
-    return src.block<3, 1>(0, index);
-  }
+  Eigen::Vector3d getBearingVector2(size_t index) const override { return src.block<3, 1>(0, index); }
 
   // TODO(nathan) think about weighted correspondences
   double getWeight(size_t index) const override { return 1.0; }
 
-  Eigen::Vector3d getCamOffset1(size_t index) const override {
-    return Eigen::Vector3d::Zero();
-  }
+  Eigen::Vector3d getCamOffset1(size_t index) const override { return Eigen::Vector3d::Zero(); }
 
-  Eigen::Matrix3d getCamRotation1(size_t index) const override {
-    return Eigen::Matrix3d::Identity();
-  }
+  Eigen::Matrix3d getCamRotation1(size_t index) const override { return Eigen::Matrix3d::Identity(); }
 
-  Eigen::Vector3d getCamOffset2(size_t index) const override {
-    return Eigen::Vector3d::Zero();
-  }
+  Eigen::Vector3d getCamOffset2(size_t index) const override { return Eigen::Vector3d::Zero(); }
 
-  Eigen::Matrix3d getCamRotation2(size_t index) const override {
-    return Eigen::Matrix3d::Identity();
-  }
+  Eigen::Matrix3d getCamRotation2(size_t index) const override { return Eigen::Matrix3d::Identity(); }
 
   size_t getNumberCorrespondences() const override { return dest.cols(); }
 
-  Vec3Matrix src;
-  Vec3Matrix dest;
+  const PyVec3Matrix src;
+  const PyVec3Matrix dest;
 };
 
 struct EigenAbsoluteAdaptor : public AbsoluteAdapterBase {
   EigenAbsoluteAdaptor(const Vec3Matrix& bearings, const Vec3Matrix& points)
       : EigenAbsoluteAdaptor(bearings, points, Eigen::Matrix3d::Identity()) {}
 
-  EigenAbsoluteAdaptor(const Vec3Matrix& bearings,
-                       const Vec3Matrix& points,
-                       const Eigen::Matrix3d& rotation_prior)
-      : EigenAbsoluteAdaptor(
-            bearings, points, rotation_prior, Eigen::Vector3d::Zero()) {}
+  EigenAbsoluteAdaptor(const Vec3Matrix& bearings, const Vec3Matrix& points, const Eigen::Matrix3d& rotation_prior)
+      : EigenAbsoluteAdaptor(bearings, points, rotation_prior, Eigen::Vector3d::Zero()) {}
 
   EigenAbsoluteAdaptor(const Vec3Matrix& _bearings,
                        const Vec3Matrix& _points,
                        const Eigen::Matrix3d& rotation_prior,
                        const Eigen::Vector3d& translation_prior)
-      : AbsoluteAdapterBase(translation_prior, rotation_prior),
-        bearings(_bearings),
-        points(_points) {
+      : AbsoluteAdapterBase(translation_prior, rotation_prior), bearings(_bearings), points(_points) {
     checkInputs(bearings, points, "bearings", "points");
   }
 
-  Eigen::Vector3d getBearingVector(size_t index) const override {
-    return bearings.block<3, 1>(0, index);
-  }
+  Eigen::Vector3d getBearingVector(size_t index) const override { return bearings.block<3, 1>(0, index); }
 
   // TODO(nathan) think about weighted correspondences
   double getWeight(size_t index) const override { return 1.0; }
 
-  Eigen::Vector3d getCamOffset(size_t index) const override {
-    return Eigen::Vector3d::Zero();
-  }
+  Eigen::Vector3d getCamOffset(size_t index) const override { return Eigen::Vector3d::Zero(); }
 
-  Eigen::Matrix3d getCamRotation(size_t index) const override {
-    return Eigen::Matrix3d::Identity();
-  }
+  Eigen::Matrix3d getCamRotation(size_t index) const override { return Eigen::Matrix3d::Identity(); }
 
-  Eigen::Vector3d getPoint(size_t index) const override {
-    return points.block<3, 1>(0, index);
-  }
+  Eigen::Vector3d getPoint(size_t index) const override { return points.block<3, 1>(0, index); }
 
   size_t getNumberCorrespondences() const override { return bearings.cols(); }
 
@@ -136,28 +101,20 @@ struct EigenPointCloudAdaptor : public PointCloudAdapterBase {
   EigenPointCloudAdaptor(const Vec3Matrix& src, const Vec3Matrix& dest)
       : EigenPointCloudAdaptor(src, dest, Eigen::Matrix3d::Identity()) {}
 
-  EigenPointCloudAdaptor(const Vec3Matrix& src,
-                         const Vec3Matrix& dest,
-                         const Eigen::Matrix3d& rotation_prior)
+  EigenPointCloudAdaptor(const Vec3Matrix& src, const Vec3Matrix& dest, const Eigen::Matrix3d& rotation_prior)
       : EigenPointCloudAdaptor(src, dest, rotation_prior, Eigen::Vector3d::Zero()) {}
 
   EigenPointCloudAdaptor(const Vec3Matrix& _src,
                          const Vec3Matrix& _dest,
                          const Eigen::Matrix3d& rotation_prior,
                          const Eigen::Vector3d& translation_prior)
-      : PointCloudAdapterBase(translation_prior, rotation_prior),
-        src(_src),
-        dest(_dest) {
+      : PointCloudAdapterBase(translation_prior, rotation_prior), src(_src), dest(_dest) {
     checkInputs(src, dest, "src_points", "dest_points");
   }
 
-  Eigen::Vector3d getPoint1(size_t index) const override {
-    return dest.block<3, 1>(0, index);
-  }
+  Eigen::Vector3d getPoint1(size_t index) const override { return dest.block<3, 1>(0, index); }
 
-  Eigen::Vector3d getPoint2(size_t index) const override {
-    return src.block<3, 1>(0, index);
-  }
+  Eigen::Vector3d getPoint2(size_t index) const override { return src.block<3, 1>(0, index); }
 
   size_t getNumberCorrespondences() const override { return dest.cols(); }
 
@@ -176,23 +133,17 @@ struct RansacResult {
 
   RansacResult() = default;
   RansacResult(const Ransac<RelativePoseProblem>& ransac)
-      : valid(true),
-        dest_T_src(Eigen::Matrix4d::Identity()),
-        inliers(ransac.inliers_.begin(), ransac.inliers_.end()) {
+      : valid(true), dest_T_src(Eigen::Matrix4d::Identity()), inliers(ransac.inliers_.begin(), ransac.inliers_.end()) {
     dest_T_src.block<3, 4>(0, 0) = ransac.model_coefficients_;
   }
 
   RansacResult(const Ransac<AbsolutePoseProblem>& ransac)
-      : valid(true),
-        dest_T_src(Eigen::Matrix4d::Identity()),
-        inliers(ransac.inliers_.begin(), ransac.inliers_.end()) {
+      : valid(true), dest_T_src(Eigen::Matrix4d::Identity()), inliers(ransac.inliers_.begin(), ransac.inliers_.end()) {
     dest_T_src.block<3, 4>(0, 0) = ransac.model_coefficients_;
   }
 
   RansacResult(const Ransac<PointCloudSacProblem>& ransac)
-      : valid(true),
-        dest_T_src(Eigen::Matrix4d::Identity()),
-        inliers(ransac.inliers_.begin(), ransac.inliers_.end()) {
+      : valid(true), dest_T_src(Eigen::Matrix4d::Identity()), inliers(ransac.inliers_.begin(), ransac.inliers_.end()) {
     dest_T_src.block<3, 4>(0, 0) = ransac.model_coefficients_;
   }
 };
@@ -205,17 +156,10 @@ PYBIND11_MODULE(_ouroboros_opengv, module) {
 
   py::class_<RansacResult>(module, "RansacResult", "Result struct for ransac result.")
       .def(
-          "__bool__",
-          [](const RansacResult& result) -> bool { return result; },
-          "Return whether or not the pose is valid.")
-      .def_readonly(
-          "valid", &RansacResult::valid, "Whether or not the pose result is valid.")
-      .def_readonly("dest_T_src",
-                    &RansacResult::dest_T_src,
-                    "Transform from the src input to the destination input.")
-      .def_readonly("inliers",
-                    &RansacResult::inliers,
-                    "Indicies of the inliers matching the model.");
+          "__bool__", [](const RansacResult& result) -> bool { return result; }, "Return validity of pose.")
+      .def_readonly("valid", &RansacResult::valid, "Whether or not the pose result is valid.")
+      .def_readonly("dest_T_src", &RansacResult::dest_T_src, "Transform from the src input to the destination input.")
+      .def_readonly("inliers", &RansacResult::inliers, "Indicies of the inliers matching the model.");
 
   py::enum_<RelativePoseProblem::Algorithm>(module, "Solver2d2d")
       .value("STEWENIUS", RelativePoseProblem::Algorithm::STEWENIUS)
@@ -224,19 +168,18 @@ PYBIND11_MODULE(_ouroboros_opengv, module) {
       .value("EIGHTPT", RelativePoseProblem::Algorithm::EIGHTPT)
       .export_values();
 
-  // NOTE(nathan) we don't expose EPNP because it spams stdout
+  // NOTE(nathan) we don't expose EPNP because it spams stdout and P2P because it uses
+  // known rotation
   py::enum_<AbsolutePoseProblem::Algorithm>(module, "Solver2d3d")
-      .value("TWOPT", AbsolutePoseProblem::TWOPT)
       .value("KNEIP", AbsolutePoseProblem::Algorithm::KNEIP)
       .value("GAO", AbsolutePoseProblem::Algorithm::GAO)
       .value("GP3P", AbsolutePoseProblem::Algorithm::GP3P)
       .export_values();
 
-  // TODO(nathan) allow for prior rotation or translation
   module.def(
       "solve_2d2d",
-      [](const Vec3Matrix& src,
-         const Vec3Matrix& dest,
+      [](const PyVec3Matrix& src,
+         const PyVec3Matrix& dest,
          RelativePoseProblem::Algorithm solver,
          size_t max_iterations,
          double threshold,
@@ -265,8 +208,8 @@ PYBIND11_MODULE(_ouroboros_opengv, module) {
 
 Returns:
   _ouroboros_opengv.RansacResult: Potentially valid dest_T_src and associated inliers)",
-      "src"_a,
-      "dest"_a,
+      "src"_a.noconvert(),
+      "dest"_a.noconvert(),
       "solver"_a = RelativePoseProblem::Algorithm::STEWENIUS,
       "max_iterations"_a = 1000,
       "threshold"_a = 1.0e-2,
@@ -314,11 +257,8 @@ PYBIND11_MODULE(_ouroboros_opengv, module) {
 
   module.def(
       "solve_3d3d",
-      [](const Vec3Matrix& src,
-         const Vec3Matrix& dest,
-         size_t max_iterations,
-         double threshold,
-         double probability) -> RansacResult {
+      [](const Vec3Matrix& src, const Vec3Matrix& dest, size_t max_iterations, double threshold, double probability)
+          -> RansacResult {
         EigenPointCloudAdaptor adaptor(src, dest);
         Ransac<PointCloudSacProblem> ransac;
         ransac.max_iterations_ = max_iterations;
@@ -342,7 +282,6 @@ PYBIND11_MODULE(_ouroboros_opengv, module) {
 
 Returns:
   _ouroboros_opengv.RansacResult: Potentially valid dest_T_src and associated inliers)",
-
       "src"_a,
       "dest"_a,
       "max_iterations"_a = 1000,

diff --git a/tests/test_opengv_interface.py b/tests/test_opengv_interface.py
@@ -16,27 +16,31 @@ def test_inverse_camera_matrix():
 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
+    features = np.array([[5.0, 2.5], [15.0, 2.5], [5.0, -2.5]])
     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
+    )
+    with np.printoptions(suppress=True):
+        print(f"expect:\n{expected}")
+        print(f"bearings:\n{bearings}")
+
     assert bearings == pytest.approx(expected)
 
 
 def _shuffle_features(features):
-    indices = np.arange(features.shape[1])
+    indices = np.arange(features.shape[0])
     np.random.shuffle(indices)
-    return indices, features[:, indices].copy()
+    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_features = np.random.normal(size=(100, 2))
     query_bearings = ogv.get_bearings(np.eye(3), query_features)
 
     yaw = np.pi / 4.0
@@ -48,10 +52,10 @@ def test_solver():
         ]
     )
     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, :]
+    match_bearings = query_bearings @ match_R_query.T + match_t_query.T
+    match_features = match_bearings[:, :2] / match_bearings[:, 2, np.newaxis]
 
-    indices = np.arange(query_bearings.shape[1])
+    indices = np.arange(query_bearings.shape[0])
     new_indices, match_features = _shuffle_features(match_features)
 
     # needs to be query -> match (so need indices that were used by shuffle for query)