Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Support mat.numpy() in Python #4356

Merged
merged 2 commits into from
Nov 18, 2022
Merged

Support mat.numpy() in Python #4356

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
70460dc
Support mat.numpy() in Python
csukuangfj Nov 18, 2022
10d44b6
apply code-format changes
csukuangfj Nov 18, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 22 additions & 104 deletions python/src/main.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,23 +251,7 @@ PYBIND11_MODULE(ncnn, m)
pybind11::pybind11_fail(ss.str());
}

size_t elemsize = 4u;
if (info.format == py::format_descriptor<double>::format())
{
elemsize = 8u;
}
if (info.format == py::format_descriptor<float>::format() || info.format == py::format_descriptor<int>::format())
{
elemsize = 4u;
}
else if (info.format == "e")
{
elemsize = 2u;
}
else if (info.format == py::format_descriptor<int8_t>::format() || info.format == py::format_descriptor<uint8_t>::format())
{
elemsize = 1u;
}
size_t elemsize = info.itemsize;

Mat* v = nullptr;
if (info.ndim == 1)
Expand Down Expand Up @@ -296,73 +280,24 @@ PYBIND11_MODULE(ncnn, m)
// so we set the cstep as numpy's cstep
v->cstep = (int)info.shape[3] * (int)info.shape[2] * (int)info.shape[1];
}
return v;
return std::unique_ptr<Mat>(v);
}),
py::arg("array"))
.def_buffer([](Mat& m) -> py::buffer_info {
if (m.elemsize != 1 && m.elemsize != 2 && m.elemsize != 4)
{
std::stringstream ss;
ss << "convert ncnn.Mat to numpy.ndarray only elemsize 1, 2, 4 support now, but given " << m.elemsize;
pybind11::pybind11_fail(ss.str());
}
if (m.elempack != 1)
{
std::stringstream ss;
ss << "convert ncnn.Mat to numpy.ndarray only elempack 1 support now, but given " << m.elempack;
pybind11::pybind11_fail(ss.str());
}
std::string format = get_mat_format(m);
std::vector<py::ssize_t> shape;
std::vector<py::ssize_t> strides;
if (m.dims == 1)
{
shape.push_back(m.w);
strides.push_back(m.elemsize);
}
else if (m.dims == 2)
{
shape.push_back(m.h);
shape.push_back(m.w);
strides.push_back(m.w * m.elemsize);
strides.push_back(m.elemsize);
}
else if (m.dims == 3)
{
shape.push_back(m.c);
shape.push_back(m.h);
shape.push_back(m.w);
strides.push_back(m.cstep * m.elemsize);
strides.push_back(m.w * m.elemsize);
strides.push_back(m.elemsize);
}
else if (m.dims == 4)
{
shape.push_back(m.c);
shape.push_back(m.d);
shape.push_back(m.h);
shape.push_back(m.w);
strides.push_back(m.cstep * m.elemsize);
strides.push_back(m.w * m.h * m.elemsize);
strides.push_back(m.w * m.elemsize);
strides.push_back(m.elemsize);
}
return py::buffer_info(
m.data, /* Pointer to buffer */
m.elemsize, /* Size of one scalar */
format, /* Python struct-style format descriptor */
m.dims, /* Number of dimensions */
shape, /* Buffer dimensions */
strides /* Strides (in bytes) for each index */
);
return to_buffer_info(m);
})
.def(
"numpy", [](py::object obj, const std::string& format = "") -> py::array {
auto* m = obj.cast<Mat*>();
return py::array(to_buffer_info(*m, format), obj);
},
py::arg("format") = "", "i for int32, f for float32, d for double")
//.def("fill", (void (Mat::*)(int))(&Mat::fill), py::arg("v"))
.def("fill", (void (Mat::*)(float))(&Mat::fill), py::arg("v"))
.def("clone", &Mat::clone, py::arg("allocator") = nullptr)
.def("clone_from", &Mat::clone_from, py::arg("mat"), py::arg("allocator") = nullptr)
.def(
"reshape",
[](Mat& mat, py::tuple shape, Allocator* allocator) {
"reshape", [](Mat& mat, py::tuple shape, Allocator* allocator) {
switch (shape.size())
{
case 1:
Expand All @@ -381,18 +316,13 @@ PYBIND11_MODULE(ncnn, m)
return Mat();
},
py::arg("shape") = py::tuple(1), py::arg("allocator") = nullptr)
.def("reshape", (Mat(Mat::*)(int, Allocator*) const) & Mat::reshape,
py::arg("w"), py::kw_only(), py::arg("allocator") = nullptr)
.def("reshape", (Mat(Mat::*)(int, int, Allocator*) const) & Mat::reshape,
py::arg("w"), py::arg("h"), py::kw_only(), py::arg("allocator") = nullptr)
.def("reshape", (Mat(Mat::*)(int, int, int, Allocator*) const) & Mat::reshape,
py::arg("w"), py::arg("h"), py::arg("c"), py::kw_only(), py::arg("allocator") = nullptr)
.def("reshape", (Mat(Mat::*)(int, int, int, int, Allocator*) const) & Mat::reshape,
py::arg("w"), py::arg("h"), py::arg("d"), py::arg("c"), py::kw_only(), py::arg("allocator") = nullptr)
.def("reshape", (Mat(Mat::*)(int, Allocator*) const) & Mat::reshape, py::arg("w"), py::kw_only(), py::arg("allocator") = nullptr)
.def("reshape", (Mat(Mat::*)(int, int, Allocator*) const) & Mat::reshape, py::arg("w"), py::arg("h"), py::kw_only(), py::arg("allocator") = nullptr)
.def("reshape", (Mat(Mat::*)(int, int, int, Allocator*) const) & Mat::reshape, py::arg("w"), py::arg("h"), py::arg("c"), py::kw_only(), py::arg("allocator") = nullptr)
.def("reshape", (Mat(Mat::*)(int, int, int, int, Allocator*) const) & Mat::reshape, py::arg("w"), py::arg("h"), py::arg("d"), py::arg("c"), py::kw_only(), py::arg("allocator") = nullptr)

.def(
"create",
[](Mat& mat, py::tuple shape, size_t elemsize, int elempack, Allocator* allocator) {
"create", [](Mat& mat, py::tuple shape, size_t elemsize, int elempack, Allocator* allocator) {
switch (shape.size())
{
case 1:
Expand All @@ -410,23 +340,12 @@ PYBIND11_MODULE(ncnn, m)
}
return;
},
py::arg("shape"), py::kw_only(),
py::arg("elemsize") = 4, py::arg("elempack") = 1,
py::arg("allocator") = nullptr)
.def("create", (void (Mat::*)(int, size_t, int, Allocator*)) & Mat::create,
py::arg("w"), py::kw_only(),
py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
.def("create", (void (Mat::*)(int, int, size_t, int, Allocator*)) & Mat::create,
py::arg("w"), py::arg("h"), py::kw_only(),
py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
.def("create", (void (Mat::*)(int, int, int, size_t, int, Allocator*)) & Mat::create,
py::arg("w"), py::arg("h"), py::arg("c"), py::kw_only(),
py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
.def("create", (void (Mat::*)(int, int, int, int, size_t, int, Allocator*)) & Mat::create,
py::arg("w"), py::arg("h"), py::arg("d"), py::arg("c"), py::kw_only(),
py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
.def("create_like", (void (Mat::*)(const Mat&, Allocator*)) & Mat::create_like,
py::arg("m"), py::arg("allocator") = nullptr)
py::arg("shape"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
.def("create", (void (Mat::*)(int, size_t, int, Allocator*)) & Mat::create, py::arg("w"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
.def("create", (void (Mat::*)(int, int, size_t, int, Allocator*)) & Mat::create, py::arg("w"), py::arg("h"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
.def("create", (void (Mat::*)(int, int, int, size_t, int, Allocator*)) & Mat::create, py::arg("w"), py::arg("h"), py::arg("c"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
.def("create", (void (Mat::*)(int, int, int, int, size_t, int, Allocator*)) & Mat::create, py::arg("w"), py::arg("h"), py::arg("d"), py::arg("c"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
.def("create_like", (void (Mat::*)(const Mat&, Allocator*)) & Mat::create_like, py::arg("m"), py::arg("allocator") = nullptr)
.def("addref", &Mat::addref)
.def("release", &Mat::release)
.def("empty", &Mat::empty)
Expand All @@ -438,8 +357,7 @@ PYBIND11_MODULE(ncnn, m)
.def("depth", (Mat(Mat::*)(int)) & Mat::depth, py::arg("z"))
//.def("depth", (const Mat (Mat::*)(int) const) & Mat::depth, py::arg("z"))
.def(
"row",
[](Mat& m, int y) {
"row", [](Mat& m, int y) {
if (m.elempack != 1)
{
std::stringstream ss;
Expand Down
74 changes: 73 additions & 1 deletion python/src/pybind11_mat.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@

#include <mat.h>

namespace py = pybind11;

std::string get_mat_format(const ncnn::Mat& m)
{
std::string format;
Expand All @@ -41,4 +43,74 @@ std::string get_mat_format(const ncnn::Mat& m)
return format;
}

#endif
// possible values for format:
// i (int32_t)
// f (float)
// d (double)
// leave it to empty to use get_mat_format
py::buffer_info to_buffer_info(ncnn::Mat& m, const std::string& format = "")
{
if (m.elemsize != 1 && m.elemsize != 2 && m.elemsize != 4)
{
std::ostringstream ss;
ss << "Convert ncnn.Mat to numpy.ndarray. Support only elemsize 1, 2, 4; but given "
<< m.elemsize;
py::pybind11_fail(ss.str());
}
if (m.elempack != 1)
{
std::ostringstream ss;
ss << "Convert ncnn.Mat to numpy.ndarray. Support only elempack == 1, but "
"given "
<< m.elempack;
py::pybind11_fail(ss.str());
}
std::string _format(format);
if (_format.empty())
{
_format = get_mat_format(m);
}
std::vector<py::ssize_t> shape;
std::vector<py::ssize_t> strides;
if (m.dims == 1)
{
shape.push_back(m.w);
strides.push_back(m.elemsize);
}
else if (m.dims == 2)
{
shape.push_back(m.h);
shape.push_back(m.w);
strides.push_back(m.w * m.elemsize);
strides.push_back(m.elemsize);
}
else if (m.dims == 3)
{
shape.push_back(m.c);
shape.push_back(m.h);
shape.push_back(m.w);
strides.push_back(m.cstep * m.elemsize);
strides.push_back(m.w * m.elemsize);
strides.push_back(m.elemsize);
}
else if (m.dims == 4)
{
shape.push_back(m.c);
shape.push_back(m.d);
shape.push_back(m.h);
shape.push_back(m.w);
strides.push_back(m.cstep * m.elemsize);
strides.push_back(m.w * m.h * m.elemsize);
strides.push_back(m.w * m.elemsize);
strides.push_back(m.elemsize);
}
return py::buffer_info(m.data, /* Pointer to buffer */
m.elemsize, /* Size of one scalar */
_format, /* Python struct-style format descriptor */
m.dims, /* Number of dimensions */
shape, /* Buffer dimensions */
strides /* Strides (in bytes) for each index */
);
}

#endif
24 changes: 19 additions & 5 deletions python/tests/test_mat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,10 +211,11 @@ def test_mat_dims4():

def test_numpy():
mat = ncnn.Mat(1)
array = np.array(mat)
array = mat.numpy()
assert mat.dims == array.ndim and mat.w == array.shape[0]
mat = ncnn.Mat(2, 3)
array = np.array(mat)
array = mat.numpy()
assert array.dtype == np.float32
assert (
mat.dims == array.ndim and mat.w == array.shape[1] and mat.h == array.shape[0]
)
Expand All @@ -237,10 +238,10 @@ def test_numpy():
)

mat = ncnn.Mat(1, elemsize=1)
array = np.array(mat)
array = mat.numpy()
assert array.dtype == np.int8
mat = ncnn.Mat(1, elemsize=2)
array = np.array(mat)
array = mat.numpy()
assert array.dtype == np.float16
# pybind11 def_buffer throw bug
# with pytest.raises(RuntimeError) as execinfo:
Expand All @@ -251,7 +252,7 @@ def test_numpy():
# )
assert array.dtype == np.float16
mat = ncnn.Mat(1, elemsize=4)
array = np.array(mat)
array = mat.numpy()
assert array.dtype == np.float32

mat = np.random.randint(0, 128, size=(12,)).astype(np.uint8)
Expand Down Expand Up @@ -279,6 +280,19 @@ def test_numpy():
array = np.array(mat)
assert (mat == array).all()

array = np.array([1, 2, 3], dtype=np.int32)
mat = ncnn.Mat(array)
array2 = mat.numpy(format='i')
assert array2.dtype == np.int32
array[0] = 10
assert array2[0] == 10

array = np.array([1, 2, 3], dtype=np.float32)
mat = ncnn.Mat(array)
array2 = mat.numpy(format='f')
assert array2.dtype == np.float32
array2[0] = 100
assert array[0] == 100

def test_fill():
mat = ncnn.Mat(1)
Expand Down