Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <pybind11/pybind11.h>
- #include <pybind11/embed.h>
- #include <pybind11/numpy.h>
- #include <pybind11/stl.h>
- #include <opencv2/opencv.hpp>
- #include <iostream>
- namespace py = pybind11;
- // ============================================================================
- int determine_cv_type(pybind11::dtype const& dtype)
- {
- switch (dtype.kind()) {
- case 'u':
- switch (dtype.itemsize()) {
- case 1: return CV_8U;
- case 2: return CV_16U;
- default:
- throw std::invalid_argument("Unsupported unsigned integer size.");
- }
- case 'i':
- switch (dtype.itemsize()) {
- case 1: return CV_8S;
- case 2: return CV_16S;
- case 4: return CV_32S;
- default:
- throw std::invalid_argument("Unsupported signed integer size.");
- }
- case 'f':
- switch (dtype.itemsize()) {
- case 4: return CV_32F;
- case 8: return CV_64F;
- default:
- throw std::invalid_argument("Unsupported float size.");
- }
- default:
- throw std::invalid_argument("Unsupported array dtype.");
- }
- }
- // ----------------------------------------------------------------------------
- cv::Mat nparray_to_mat(py::array& arr)
- {
- auto const depth = determine_cv_type(arr.dtype());
- switch (arr.ndim()) {
- case 1:
- return cv::Mat(
- static_cast<int>(arr.shape(0)) // Rows
- , 1 // Columns
- , CV_MAKETYPE(depth, 1) // Data type
- , arr.mutable_data()
- );
- case 2:
- return cv::Mat(
- static_cast<int>(arr.shape(0)) // Rows
- , static_cast<int>(arr.shape(1)) // Columns
- , CV_MAKETYPE(depth, 1) // Data type
- , arr.mutable_data()
- );
- case 3:
- if (arr.shape(2) > CV_CN_MAX) {
- std::invalid_argument("Too many array channels.");
- }
- return cv::Mat(
- static_cast<int>(arr.shape(0)) // Rows
- , static_cast<int>(arr.shape(1)) // Columns
- , CV_MAKETYPE(depth, static_cast<int>(arr.shape(2))) // Data type
- , arr.mutable_data()
- );
- default:
- throw std::invalid_argument("Only 1D, 2D and 3D arrays supported.");
- }
- }
- // ============================================================================
- py::dtype determine_np_dtype(int depth)
- {
- switch (depth) {
- case CV_8U: return py::dtype::of<uint8_t>();
- case CV_8S: return py::dtype::of<int8_t>();
- case CV_16U: return py::dtype::of<uint16_t>();
- case CV_16S: return py::dtype::of<int16_t>();
- case CV_32S: return py::dtype::of<int32_t>();
- case CV_32F: return py::dtype::of<float>();
- case CV_64F: return py::dtype::of<double>();
- default:
- throw std::invalid_argument("Unsupported data type.");
- }
- }
- // ----------------------------------------------------------------------------
- py::capsule make_capsule(cv::Mat& m)
- {
- return py::capsule(new cv::Mat(m)
- , [](void *v) { delete reinterpret_cast<cv::Mat*>(v); }
- );
- }
- // ----------------------------------------------------------------------------
- py::array mat_to_nparray(cv::Mat& m)
- {
- if (!m.isContinuous()) {
- throw std::invalid_argument("Only continuous Mats supported.");
- }
- std::vector<std::size_t> shape;
- if (m.channels() == 1) {
- shape = {
- static_cast<size_t>(m.rows)
- , static_cast<size_t>(m.cols)
- };
- } else {
- shape = {
- static_cast<size_t>(m.rows)
- , static_cast<size_t>(m.cols)
- , static_cast<size_t>(m.channels())
- };
- }
- return py::array(determine_np_dtype(m.depth())
- , shape
- , m.data
- , make_capsule(m));
- }
- // ============================================================================
- PYBIND11_EMBEDDED_MODULE(test_module, m)
- {
- m.doc() = "Test module";
- }
- // ============================================================================
- int main()
- {
- // Start the interpreter and keep it alive
- py::scoped_interpreter guard{};
- try {
- auto locals = py::dict{};
- py::exec(R"(
- import numpy as np
- def test_cpp_to_py(arr):
- return (1,2,3)
- )");
- auto test_cpp_to_py = py::globals()["test_cpp_to_py"];
- for (int i = 0; i < 10; i++) {
- int64 t0 = cv::getTickCount();
- cv::Mat img(cv::Mat::zeros(1024, 1024, CV_8UC3) + cv::Scalar(1, 1, 1));
- int64 t1 = cv::getTickCount();
- auto result = test_cpp_to_py(mat_to_nparray(img));
- int64 t2 = cv::getTickCount();
- double delta0 = (t1 - t0) / cv::getTickFrequency() * 1000;
- double delta1 = (t2 - t1) / cv::getTickFrequency() * 1000;
- std::cout << "* " << delta0 << " ms | " << delta1 << " ms" << std::endl;
- }
- } catch (py::error_already_set& e) {
- std::cerr << e.what() << "\n";
- }
- return 0;
- }
- // ============================================================================
Add Comment
Please, Sign In to add comment