pybind11 - cv::Mat <-> np.array

#include <pybind11/pybind11.h>
#include <pybind11/embed.h>
#include <pybind11/numpy.h>

#include <opencv2/opencv.hpp>

#include <iostream>

namespace py = pybind11;

// Restrict to only C-Style memory layout
void process_array(py::array_t<uint8_t, py::array::c_style>& arr)
{
    if (arr.ndim() != 3) {
        throw std::runtime_error("Incorrect array shape.");
    }

    // Note: Shares buffer with the Python array -- limit to current function
    // or make a deep copy (otherwise assure the source array lifetime is longer)
    cv::Mat img(arr.shape(0), arr.shape(1)
        , CV_8UC3
        , static_cast<uchar*>(arr.mutable_data()));

    std::cout << "Rows = " << img.rows << "\n";
    std::cout << "Columns = " << img.cols << "\n";
    std::cout << "Channels = " << img.channels() << "\n";
    std::cout << img << "\n";
}

PYBIND11_EMBEDDED_MODULE(test_module, m)
{
    m.doc() = "Test module";

    m.def("process_array", &process_array, "Test function");
}

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
            import test_module

            def test_py_to_cpp():
                a = np.arange(4*5*3, dtype=np.uint8).reshape(4,5,3)
                test_module.process_array(a)

            def test_cpp_to_py(arr):
                arr += 1
        )");

        //auto test_py_to_cpp = py::globals()["test_py_to_cpp"];
        //test_py_to_cpp();

        auto test_cpp_to_py = py::globals()["test_cpp_to_py"];

        cv::Mat img(cv::Mat::zeros(4, 5, CV_8UC3) + cv::Scalar(1,1,1));

        // Create a numpy array that shares the buffer with the Mat
        /*
        py::array_t<uint8_t, py::array::c_style> arr(
            py::buffer_info(
                img.data // Buffer
                , sizeof(uint8_t) // Item size
                , py::format_descriptor<uint8_t>::format()
                , 3 // Dimension count
                , std::vector<size_t> { static_cast<size_t>(img.rows)
                    , static_cast<size_t>(img.cols)
                    , static_cast<size_t>(img.channels()) } // Dimensions
                , std::vector<size_t> { sizeof(uint8_t) * img.channels() * img.cols
                    , sizeof(uint8_t) * img.channels()
                    , sizeof(uint8_t) } // Strides
            ));
        */

        py::array_t<uint8_t, py::array::c_style> arr(
            // Dimensions of the array
            {   static_cast<size_t>(img.rows)
                , static_cast<size_t>(img.cols)
                , static_cast<size_t>(img.channels())
            }
            // Data buffer
            , img.data
            // Holds an extra reference to the Mat owning the shared buffer
            , py::capsule(new cv::Mat(img), [](void *m) { delete reinterpret_cast<cv::Mat*>(m); })
        );

        std::cout << "Before\n" << img << "\n";
        test_cpp_to_py(arr);
        std::cout << "After\n" << img << "\n";
    } catch (py::error_already_set& e) {
        std::cerr << e.what() << "\n";
    }

    return 0;
}