GithubStar/ncnn
1
0
Fork 0
mirror of https://github.com/Tencent/ncnn.git synced 2026-10-23 09:08:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Support mat.numpy() in Python (#4356)

Browse Source
This commit is contained in:
Fangjun Kuang
2022-11-18 15:47:02 +08:00
committed by GitHub
parent 057b5bb515
commit 1b83fe4f16
3 changed files with 114 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/main.cpp
+22 -104
View File
@@ -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)
@@ -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:
@@ -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:
@@ -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)
@@ -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;