Autocast numpy.ndarray <-> std::vector<Eigen::Vector>

pycolmap/main.cc

@@ -12,7 +12,6 @@
 #include "pycolmap/utils.h"
 
 #include <glog/logging.h>
-#include <pybind11/eigen.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 

pycolmap/pybind11_extension.h

@@ -6,6 +6,7 @@
 #include <string>
 
 #include <pybind11/cast.h>
+#include <pybind11/eigen.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/pytypes.h>
 #include <pybind11/stl.h>
@@ -56,6 +57,43 @@ struct type_caster<std::string> {
   }
 };
 
+// Autocast from numpy.ndarray to std::vector<Eigen::Vector>
+template <typename Scalar, int Size>
+struct type_caster<std::vector<Eigen::Matrix<Scalar, Size, 1>>> {
+ public:
+  using MatrixType =
+      typename Eigen::Matrix<Scalar, Eigen::Dynamic, Size, Eigen::RowMajor>;
+  using VectorType = typename Eigen::Matrix<Scalar, Size, 1>;
+  using props = EigenProps<MatrixType>;
+  PYBIND11_TYPE_CASTER(std::vector<VectorType>, props::descriptor);
+
+  bool load(handle src, bool) {
+    const auto buf = array::ensure(src);
+    if (!buf) {
+      return false;
+    }
+    const buffer_info info = buf.request();
+    if (info.ndim != 2 || info.shape[1] != Size) {
+      return false;
+    }
+    const size_t num_elements = info.shape[0];
+    value.resize(num_elements);
+    const auto& mat = src.cast<Eigen::Ref<const MatrixType>>();
+    Eigen::Map<MatrixType>(
+        reinterpret_cast<Scalar*>(value.data()), num_elements, Size) = mat;
+    return true;
+  }
+
+  static handle cast(const std::vector<VectorType>& vec,
+                     return_value_policy /* policy */,
+                     handle h) {
+    Eigen::Map<const MatrixType> mat(
+        reinterpret_cast<const Scalar*>(vec.data()), vec.size(), Size);
+    return type_caster<Eigen::Map<const MatrixType>>::cast(
+        mat, return_value_policy::copy, h);
+  }
+};
+
 }  // namespace detail
 
 // Fix long-standing bug https://github.com/pybind/pybind11/issues/4529

pycolmap/scene/camera.h

@@ -140,22 +140,23 @@ void BindCamera(py::module& m) {
            "Project point in image plane to world / infinity.")
       .def(
           "cam_from_img",
-          [](const Camera& self, const std::vector<Eigen::Vector2d>& points2D) {
-            std::vector<Eigen::Vector2d> world_points2D;
-            for (size_t idx = 0; idx < points2D.size(); ++idx) {
-              world_points2D.push_back(self.CamFromImg(points2D[idx]));
+          [](const Camera& self,
+             const py::EigenDRef<const Eigen::MatrixX2d>& image_points) {
+            std::vector<Eigen::Vector2d> world_points(image_points.rows());
+            for (size_t idx = 0; idx < image_points.rows(); ++idx) {
+              world_points[idx] = self.CamFromImg(image_points.row(idx));
             }
-            return world_points2D;
+            return world_points;
           },
           "Project list of points in image plane to world / infinity.")
       .def(
           "cam_from_img",
-          [](const Camera& self, const std::vector<Point2D>& points2D) {
-            std::vector<Eigen::Vector2d> world_points2D;
-            for (size_t idx = 0; idx < points2D.size(); ++idx) {
-              world_points2D.push_back(self.CamFromImg(points2D[idx].xy));
+          [](const Camera& self, const std::vector<Point2D>& image_points) {
+            std::vector<Eigen::Vector2d> world_points(image_points.size());
+            for (size_t idx = 0; idx < image_points.size(); ++idx) {
+              world_points[idx] = self.CamFromImg(image_points[idx].xy);
             }
-            return world_points2D;
+            return world_points;
           },
           "Project list of points in image plane to world / infinity.")
       .def("cam_from_img_threshold",
@@ -167,22 +168,22 @@ void BindCamera(py::module& m) {
       .def(
           "img_from_cam",
           [](const Camera& self,
-             const std::vector<Eigen::Vector2d>& world_points2D) {
-            std::vector<Eigen::Vector2d> image_points2D;
-            for (size_t idx = 0; idx < world_points2D.size(); ++idx) {
-              image_points2D.push_back(self.ImgFromCam(world_points2D[idx]));
+             const py::EigenDRef<const Eigen::MatrixX2d>& world_points) {
+            std::vector<Eigen::Vector2d> image_points(world_points.rows());
+            for (size_t idx = 0; idx < world_points.rows(); ++idx) {
+              image_points[idx] = self.ImgFromCam(world_points.row(idx));
             }
-            return image_points2D;
+            return image_points;
           },
           "Project list of points from world / infinity to image plane.")
       .def(
           "img_from_cam",
-          [](const Camera& self, const std::vector<Point2D>& world_points2D) {
-            std::vector<Eigen::Vector2d> image_points2D;
-            for (size_t idx = 0; idx < world_points2D.size(); ++idx) {
-              image_points2D.push_back(self.ImgFromCam(world_points2D[idx].xy));
+          [](const Camera& self, const std::vector<Point2D>& world_points) {
+            std::vector<Eigen::Vector2d> image_points(world_points.size());
+            for (size_t idx = 0; idx < world_points.size(); ++idx) {
+              image_points[idx] = self.ImgFromCam(world_points[idx].xy);
             }
-            return image_points2D;
+            return image_points;
           },
           "Project list of points from world / infinity to image plane.")
       .def("rescale",