From 4d4c10a2dc017e25343a5703b1037d3d45fa3bff Mon Sep 17 00:00:00 2001
From: Li Zhang <lzhang329@gmail.com>
Date: Thu, 24 Nov 2022 18:23:34 +0800
Subject: [PATCH] [Enhancement] Avoid copying dense arrays in Python API
 (#1349)

* eliminate copying for segmentor

* fix segmentor

* eliminate copying in Python API

* minor fix
---
 csrc/mmdeploy/apis/python/detector.cpp  | 14 +++++++++-----
 csrc/mmdeploy/apis/python/restorer.cpp  | 20 ++++++++++++--------
 csrc/mmdeploy/apis/python/segmentor.cpp | 22 +++++++++++++---------
 3 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/csrc/mmdeploy/apis/python/detector.cpp b/csrc/mmdeploy/apis/python/detector.cpp
index 862132afc..057a92ab0 100644
--- a/csrc/mmdeploy/apis/python/detector.cpp
+++ b/csrc/mmdeploy/apis/python/detector.cpp
@@ -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() {
diff --git a/csrc/mmdeploy/apis/python/restorer.cpp b/csrc/mmdeploy/apis/python/restorer.cpp
index e4ec5ef83..771af2a6c 100644
--- a/csrc/mmdeploy/apis/python/restorer.cpp
+++ b/csrc/mmdeploy/apis/python/restorer.cpp
@@ -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:
diff --git a/csrc/mmdeploy/apis/python/segmentor.cpp b/csrc/mmdeploy/apis/python/segmentor.cpp
index e148ceda1..1fdf719fc 100644
--- a/csrc/mmdeploy/apis/python/segmentor.cpp
+++ b/csrc/mmdeploy/apis/python/segmentor.cpp
@@ -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: