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

* eliminate copying for segmentor

* fix segmentor

* eliminate copying in Python API

* minor fix
pull/1366/head
Li Zhang 2022-11-24 18:23:34 +08:00 committed by GitHub
parent 73e095a4b8
commit 4d4c10a2dc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 34 additions and 22 deletions

View File

@ -28,12 +28,16 @@ class PyDetector {
if (status != MMDEPLOY_SUCCESS) {
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 result = detection;
for (int i = 0; i < mats.size(); ++i) {
auto bboxes = py::array_t<float>({result_count[i], 5});
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]);
for (int j = 0; j < result_count[i]; ++j, ++result) {
auto bbox = bboxes.mutable_data(j);
@ -44,16 +48,16 @@ class PyDetector {
bbox[4] = result->score;
labels.mutable_at(j) = result->label_id;
if (result->mask) {
py::array_t<uint8_t> mask({result->mask->height, result->mask->width});
memcpy(mask.mutable_data(), result->mask->data, mask.nbytes());
masks.push_back(std::move(mask));
masks.emplace_back(std::array{result->mask->height, result->mask->width}, // shape
reinterpret_cast<uint8_t*>(result->mask->data), // data
py::capsule(new Sptr(holder), // handle
[](void* p) { delete reinterpret_cast<Sptr*>(p); }));
} else {
masks.emplace_back();
}
}
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;
}
~PyDetector() {

View File

@ -19,7 +19,7 @@ class PyRestorer {
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;
mats.reserve(imgs.size());
for (const auto& img : imgs) {
@ -31,15 +31,19 @@ class PyRestorer {
if (status != MMDEPLOY_SUCCESS) {
throw std::runtime_error("failed to apply restorer, code: " + std::to_string(status));
}
auto output = std::vector<py::array_t<uint8_t>>{};
output.reserve(mats.size());
using Sptr = std::shared_ptr<mmdeploy_mat_t>;
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) {
py::array_t<uint8_t> restored({results[i].height, results[i].width, results[i].channel});
memcpy(restored.mutable_data(), results[i].data, restored.nbytes());
output.push_back(std::move(restored));
rets[i] = {
{results[i].height, results[i].width, results[i].channel}, // shape
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 output;
return rets;
}
private:

View File

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