Autocast numpy.ndarray <-> std::vector<Eigen::Vector> (#248)

100x speedup for 10k inputs.

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",