[Enhancement] Avoid copying dense arrays in Python API (#1349)

* eliminate copying for segmentor

* fix segmentor

* eliminate copying in Python API

* minor fix
(cherry picked from commit 4d4c10a2dc017e25343a5703b1037d3d45fa3bff)
This commit is contained in:
Li Zhang 2022-11-24 18:23:34 +08:00 committed by lvhan028
parent a1ca93ea1c
commit da96684128
3 changed files with 34 additions and 22 deletions

View File

@ -28,12 +28,16 @@ class PyDetector {
if (status != MMDEPLOY_SUCCESS) { if (status != MMDEPLOY_SUCCESS) {
throw std::runtime_error("failed to apply detector, code: " + std::to_string(status)); throw std::runtime_error("failed to apply detector, code: " + std::to_string(status));
} }
using Sptr = std::shared_ptr<mmdeploy_detection_t>;
Sptr holder(detection, [result_count, n = mats.size()](auto p) {
mmdeploy_detector_release_result(p, result_count, n);
});
auto output = py::list{}; auto output = py::list{};
auto result = detection; auto result = detection;
for (int i = 0; i < mats.size(); ++i) { for (int i = 0; i < mats.size(); ++i) {
auto bboxes = py::array_t<float>({result_count[i], 5}); auto bboxes = py::array_t<float>({result_count[i], 5});
auto labels = py::array_t<int>(result_count[i]); auto labels = py::array_t<int>(result_count[i]);
auto masks = std::vector<py::array_t<uint8_t>>{}; auto masks = std::vector<py::array>();
masks.reserve(result_count[i]); masks.reserve(result_count[i]);
for (int j = 0; j < result_count[i]; ++j, ++result) { for (int j = 0; j < result_count[i]; ++j, ++result) {
auto bbox = bboxes.mutable_data(j); auto bbox = bboxes.mutable_data(j);
@ -44,16 +48,16 @@ class PyDetector {
bbox[4] = result->score; bbox[4] = result->score;
labels.mutable_at(j) = result->label_id; labels.mutable_at(j) = result->label_id;
if (result->mask) { if (result->mask) {
py::array_t<uint8_t> mask({result->mask->height, result->mask->width}); masks.emplace_back(std::array{result->mask->height, result->mask->width}, // shape
memcpy(mask.mutable_data(), result->mask->data, mask.nbytes()); reinterpret_cast<uint8_t*>(result->mask->data), // data
masks.push_back(std::move(mask)); py::capsule(new Sptr(holder), // handle
[](void* p) { delete reinterpret_cast<Sptr*>(p); }));
} else { } else {
masks.emplace_back(); masks.emplace_back();
} }
} }
output.append(py::make_tuple(std::move(bboxes), std::move(labels), std::move(masks))); output.append(py::make_tuple(std::move(bboxes), std::move(labels), std::move(masks)));
} }
mmdeploy_detector_release_result(detection, result_count, (int)mats.size());
return output; return output;
} }
~PyDetector() { ~PyDetector() {

View File

@ -19,7 +19,7 @@ class PyRestorer {
restorer_ = {}; restorer_ = {};
} }
std::vector<py::array_t<uint8_t>> Apply(const std::vector<PyImage>& imgs) { std::vector<py::array> Apply(const std::vector<PyImage>& imgs) {
std::vector<mmdeploy_mat_t> mats; std::vector<mmdeploy_mat_t> mats;
mats.reserve(imgs.size()); mats.reserve(imgs.size());
for (const auto& img : imgs) { for (const auto& img : imgs) {
@ -31,15 +31,19 @@ class PyRestorer {
if (status != MMDEPLOY_SUCCESS) { if (status != MMDEPLOY_SUCCESS) {
throw std::runtime_error("failed to apply restorer, code: " + std::to_string(status)); throw std::runtime_error("failed to apply restorer, code: " + std::to_string(status));
} }
auto output = std::vector<py::array_t<uint8_t>>{}; using Sptr = std::shared_ptr<mmdeploy_mat_t>;
output.reserve(mats.size()); Sptr holder(results, [n = mats.size()](auto p) { mmdeploy_restorer_release_result(p, n); });
std::vector<py::array> rets(mats.size());
for (int i = 0; i < mats.size(); ++i) { for (int i = 0; i < mats.size(); ++i) {
py::array_t<uint8_t> restored({results[i].height, results[i].width, results[i].channel}); rets[i] = {
memcpy(restored.mutable_data(), results[i].data, restored.nbytes()); {results[i].height, results[i].width, results[i].channel}, // shape
output.push_back(std::move(restored)); results[i].data, // data
py::capsule(new Sptr(holder), // handle
[](void* p) { delete reinterpret_cast<Sptr*>(p); }) //
};
} }
mmdeploy_restorer_release_result(results, (int)mats.size()); return rets;
return output;
} }
private: private:

View File

@ -20,7 +20,7 @@ class PySegmentor {
segmentor_ = {}; segmentor_ = {};
} }
std::vector<py::array_t<int>> Apply(const std::vector<PyImage>& imgs) { std::vector<py::array> Apply(const std::vector<PyImage>& imgs) {
std::vector<mmdeploy_mat_t> mats; std::vector<mmdeploy_mat_t> mats;
mats.reserve(imgs.size()); mats.reserve(imgs.size());
for (const auto& img : imgs) { for (const auto& img : imgs) {
@ -32,15 +32,19 @@ class PySegmentor {
if (status != MMDEPLOY_SUCCESS) { if (status != MMDEPLOY_SUCCESS) {
throw std::runtime_error("failed to apply segmentor, code: " + std::to_string(status)); throw std::runtime_error("failed to apply segmentor, code: " + std::to_string(status));
} }
auto output = std::vector<py::array_t<int>>{}; using Sptr = std::shared_ptr<mmdeploy_segmentation_t>;
output.reserve(mats.size()); Sptr holder(segm, [n = mats.size()](auto p) { mmdeploy_segmentor_release_result(p, n); });
for (int i = 0; i < mats.size(); ++i) {
auto mask = py::array_t<int>({segm[i].height, segm[i].width}); std::vector<py::array> rets(mats.size());
memcpy(mask.mutable_data(), segm[i].mask, mask.nbytes()); for (size_t i = 0; i < mats.size(); ++i) {
output.push_back(std::move(mask)); rets[i] = {
{segm[i].height, segm[i].width}, // shape
segm[i].mask, // data
py::capsule(new Sptr(holder), // handle
[](void* p) { delete reinterpret_cast<Sptr*>(p); }) //
};
} }
mmdeploy_segmentor_release_result(segm, (int)mats.size()); return rets;
return output;
} }
private: private: