diff --git a/pycolmap/main.cc b/pycolmap/main.cc index 2c20203..e152aef 100644 --- a/pycolmap/main.cc +++ b/pycolmap/main.cc @@ -12,7 +12,6 @@ #include "pycolmap/utils.h" #include -#include #include #include diff --git a/pycolmap/pybind11_extension.h b/pycolmap/pybind11_extension.h index 312d8ef..a96d0c9 100644 --- a/pycolmap/pybind11_extension.h +++ b/pycolmap/pybind11_extension.h @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -56,6 +57,43 @@ struct type_caster { } }; +// Autocast from numpy.ndarray to std::vector +template +struct type_caster>> { + public: + using MatrixType = + typename Eigen::Matrix; + using VectorType = typename Eigen::Matrix; + using props = EigenProps; + PYBIND11_TYPE_CASTER(std::vector, 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::Map( + reinterpret_cast(value.data()), num_elements, Size) = mat; + return true; + } + + static handle cast(const std::vector& vec, + return_value_policy /* policy */, + handle h) { + Eigen::Map mat( + reinterpret_cast(vec.data()), vec.size(), Size); + return type_caster>::cast( + mat, return_value_policy::copy, h); + } +}; + } // namespace detail // Fix long-standing bug https://github.com/pybind/pybind11/issues/4529