open-mmlab · lvhan028 · Dec 6, 2022 · Nov 15, 2022 · Nov 22, 2022 · Nov 24, 2022
diff --git a/csrc/mmdeploy/apis/c/mmdeploy/executor.cpp b/csrc/mmdeploy/apis/c/mmdeploy/executor.cpp
@@ -13,10 +13,10 @@ namespace {
 
 mmdeploy_scheduler_t CreateScheduler(const char* type, const Value& config = Value()) {
   try {
-    auto creator = Registry<SchedulerType>::Get().GetCreator(type);
+    auto creator = gRegistry<SchedulerType>().Get(type);
     if (!creator) {
       MMDEPLOY_ERROR("Creator for {} not found. Available schedulers: {}", type,
-                     Registry<SchedulerType>::Get().List());
+                     gRegistry<SchedulerType>().List());
       return nullptr;
     }
     return Cast(new SchedulerType(creator->Create(config)));

diff --git 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
@@ -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
@@ -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:

diff --git a/csrc/mmdeploy/backend_ops/onnxruntime/grid_sample/grid_sample.cpp b/csrc/mmdeploy/backend_ops/onnxruntime/grid_sample/grid_sample.cpp
@@ -13,8 +13,8 @@ namespace mmdeploy {
 #define MAX(a, b) (((a) < (b)) ? (b) : (a))
 #define CLIP_COORDINATES(in, out, clip_limit) out = MIN((clip_limit - 1), MAX(in, 0))
 
-GridSampleKernel::GridSampleKernel(OrtApi api, const OrtKernelInfo *info)
-    : api_(api), ort_(api_), info_(info) {
+GridSampleKernel::GridSampleKernel(const OrtApi &api, const OrtKernelInfo *info)
+    : ort_(api), info_(info) {
   align_corners_ = ort_.KernelInfoGetAttribute<int64_t>(info, "align_corners");
   interpolation_mode_ = ort_.KernelInfoGetAttribute<int64_t>(info, "interpolation_mode");
   padding_mode_ = ort_.KernelInfoGetAttribute<int64_t>(info, "padding_mode");

diff --git a/csrc/mmdeploy/backend_ops/onnxruntime/grid_sample/grid_sample.h b/csrc/mmdeploy/backend_ops/onnxruntime/grid_sample/grid_sample.h
@@ -7,12 +7,11 @@
 namespace mmdeploy {
 
 struct GridSampleKernel {
-  GridSampleKernel(OrtApi api, const OrtKernelInfo *info);
+  GridSampleKernel(const OrtApi &api, const OrtKernelInfo *info);
 
   void Compute(OrtKernelContext *context);
 
  protected:
-  OrtApi api_;
   Ort::CustomOpApi ort_;
   const OrtKernelInfo *info_;
   Ort::AllocatorWithDefaultOptions allocator_;
@@ -23,7 +22,7 @@ struct GridSampleKernel {
 };
 
 struct GridSampleOp : Ort::CustomOpBase<GridSampleOp, GridSampleKernel> {
-  void *CreateKernel(OrtApi api, const OrtKernelInfo *info) const {
+  void *CreateKernel(const OrtApi &api, const OrtKernelInfo *info) const {
     return new GridSampleKernel(api, info);
   };
 

diff --git a/csrc/mmdeploy/backend_ops/onnxruntime/modulated_deform_conv/modulated_deform_conv.cpp b/csrc/mmdeploy/backend_ops/onnxruntime/modulated_deform_conv/modulated_deform_conv.cpp
@@ -109,8 +109,9 @@ void deformable_conv2d_ref_fp32(const float *src, const float *offset, const flo
   }
 }
 
-MMCVModulatedDeformConvKernel::MMCVModulatedDeformConvKernel(OrtApi api, const OrtKernelInfo *info)
-    : api_(api), ort_(api_), info_(info) {
+MMCVModulatedDeformConvKernel::MMCVModulatedDeformConvKernel(const OrtApi &api,
+                                                             const OrtKernelInfo *info)
+    : ort_(api), info_(info) {
   std::vector<int64_t> stride = ort_.KernelInfoGetAttribute<std::vector<int64_t>>(info, "stride");
   stride_height_ = stride[0];
   stride_width_ = stride[1];

diff --git a/csrc/mmdeploy/backend_ops/onnxruntime/modulated_deform_conv/modulated_deform_conv.h b/csrc/mmdeploy/backend_ops/onnxruntime/modulated_deform_conv/modulated_deform_conv.h
@@ -7,12 +7,11 @@
 namespace mmdeploy {
 
 struct MMCVModulatedDeformConvKernel {
-  MMCVModulatedDeformConvKernel(OrtApi api, const OrtKernelInfo *info);
+  MMCVModulatedDeformConvKernel(const OrtApi &api, const OrtKernelInfo *info);
 
   void Compute(OrtKernelContext *context);
 
  protected:
-  OrtApi api_;
   Ort::CustomOpApi ort_;
   const OrtKernelInfo *info_;
   Ort::AllocatorWithDefaultOptions allocator_;
@@ -29,7 +28,7 @@ struct MMCVModulatedDeformConvKernel {
 
 struct MMCVModulatedDeformConvOp
     : Ort::CustomOpBase<MMCVModulatedDeformConvOp, MMCVModulatedDeformConvKernel> {
-  void *CreateKernel(OrtApi api, const OrtKernelInfo *info) const {
+  void *CreateKernel(const OrtApi &api, const OrtKernelInfo *info) const {
     return new MMCVModulatedDeformConvKernel(api, info);
   }
 

diff --git a/csrc/mmdeploy/backend_ops/onnxruntime/nms_rotated/nms_rotated.cpp b/csrc/mmdeploy/backend_ops/onnxruntime/nms_rotated/nms_rotated.cpp
@@ -261,8 +261,8 @@ float rotated_boxes_intersection(const RotatedBox& box1, const RotatedBox& box2)
   return polygon_area(orderedPts, num_convex);
 }
 
-NMSRotatedKernel::NMSRotatedKernel(OrtApi api, const OrtKernelInfo* info)
-    : api_(api), ort_(api_), info_(info) {
+NMSRotatedKernel::NMSRotatedKernel(const OrtApi& api, const OrtKernelInfo* info)
+    : ort_(api), info_(info) {
   iou_threshold_ = ort_.KernelInfoGetAttribute<float>(info, "iou_threshold");
   score_threshold_ = ort_.KernelInfoGetAttribute<float>(info, "score_threshold");
 

diff --git a/csrc/mmdeploy/backend_ops/onnxruntime/nms_rotated/nms_rotated.h b/csrc/mmdeploy/backend_ops/onnxruntime/nms_rotated/nms_rotated.h
@@ -12,12 +12,11 @@
 
 namespace mmdeploy {
 struct NMSRotatedKernel {
-  NMSRotatedKernel(OrtApi api, const OrtKernelInfo* info);
+  NMSRotatedKernel(const OrtApi& api, const OrtKernelInfo* info);
 
   void Compute(OrtKernelContext* context);
 
  private:
-  OrtApi api_;
   Ort::CustomOpApi ort_;
   const OrtKernelInfo* info_;
   Ort::AllocatorWithDefaultOptions allocator_;
@@ -26,7 +25,7 @@ struct NMSRotatedKernel {
 };
 
 struct NMSRotatedOp : Ort::CustomOpBase<NMSRotatedOp, NMSRotatedKernel> {
-  void* CreateKernel(OrtApi api, const OrtKernelInfo* info) const {
+  void* CreateKernel(const OrtApi& api, const OrtKernelInfo* info) const {
     return new NMSRotatedKernel(api, info);
   }
   const char* GetName() const { return "NMSRotated"; }

diff --git a/csrc/mmdeploy/codebase/common.h b/csrc/mmdeploy/codebase/common.h
@@ -32,8 +32,7 @@ class Context {
 template <class Tag>
 class CodebaseCreator : public Creator<Module> {
  public:
-  const char* GetName() const override { return Tag::name; }
-  int GetVersion() const override { return 1; }
+  std::string_view name() const noexcept override { return Tag::name; }
   std::unique_ptr<Module> Create(const Value& cfg) override {
     constexpr auto key{"component"};
     if (!cfg.contains(key)) {
@@ -45,36 +44,34 @@ class CodebaseCreator : public Creator<Module> {
       throw_exception(eInvalidArgument);
     }
     auto postprocess_type = cfg[key].get<std::string>();
-    auto creator = Registry<Tag>::Get().GetCreator(postprocess_type);
+    auto creator = gRegistry<Tag>().Get(postprocess_type);
     if (creator == nullptr) {
       MMDEPLOY_ERROR("Could not found entry '{}' in {}. Available components: {}", postprocess_type,
-                     Tag::name, Registry<Tag>::Get().List());
+                     Tag::name, gRegistry<Tag>().List());
       throw_exception(eEntryNotFound);
     }
     return creator->Create(cfg);
   }
 };
 
-#define DECLARE_CODEBASE(codebase_type, codebase_name)               \
+#define MMDEPLOY_DECLARE_CODEBASE(codebase_type, codebase_name)      \
   class codebase_type : public Context {                             \
    public:                                                           \
     static constexpr const auto name = #codebase_name;               \
     using type = std::unique_ptr<Module>;                            \
     explicit codebase_type(const Value& config) : Context(config) {} \
-  };
+  };                                                                 \
+  MMDEPLOY_DECLARE_REGISTRY(codebase_type, std::unique_ptr<Module>(const Value& config));
 
-#define REGISTER_CODEBASE(codebase)                       \
+#define MMDEPLOY_REGISTER_CODEBASE(codebase)              \
   using codebase##_##Creator = CodebaseCreator<codebase>; \
-  REGISTER_MODULE(Module, codebase##_##Creator)
+  MMDEPLOY_REGISTER_CREATOR(Module, codebase##_##Creator) \
+  MMDEPLOY_DEFINE_REGISTRY(codebase)
 
-#define REGISTER_CODEBASE_COMPONENT(codebase, component_type)                                      \
-  class component_type##_##Creator : public Creator<codebase> {                                    \
-   public:                                                                                         \
-    const char* GetName() const override { return #component_type; }                               \
-    int GetVersion() const override { return 1; }                                                  \
-    ReturnType Create(const Value& config) override { return CreateTask(component_type(config)); } \
-  };                                                                                               \
-  REGISTER_MODULE(codebase, component_type##_##Creator)
+#define MMDEPLOY_REGISTER_CODEBASE_COMPONENT(codebase, component_type)                    \
+  MMDEPLOY_REGISTER_FACTORY_FUNC(codebase, (component_type, 0), [](const Value& config) { \
+    return CreateTask(component_type(config));                                            \
+  })
 
 }  // namespace mmdeploy
 

diff --git a/csrc/mmdeploy/codebase/mmaction/base_head.cpp b/csrc/mmdeploy/codebase/mmaction/base_head.cpp
@@ -61,10 +61,12 @@ class BaseHead : public MMAction {
   int topk_{1};
 };
 
-REGISTER_CODEBASE_COMPONENT(MMAction, BaseHead);
+MMDEPLOY_REGISTER_CODEBASE_COMPONENT(MMAction, BaseHead);
+
 using SlowFastHead = BaseHead;
-REGISTER_CODEBASE_COMPONENT(MMAction, SlowFastHead);
+MMDEPLOY_REGISTER_CODEBASE_COMPONENT(MMAction, SlowFastHead);
+
 using TSNHead = BaseHead;
-REGISTER_CODEBASE_COMPONENT(MMAction, TSNHead);
+MMDEPLOY_REGISTER_CODEBASE_COMPONENT(MMAction, TSNHead);
 
 }  // namespace mmdeploy::mmaction
diff --git a/csrc/mmdeploy/codebase/mmaction/cpu/format_shape_impl.cpp b/csrc/mmdeploy/codebase/mmaction/cpu/format_shape_impl.cpp
@@ -123,16 +123,7 @@ class FormatShapeImpl : public ::mmdeploy::FormatShapeImpl {
   constexpr static Device kHost{0, 0};
 };
 
-class FormatShapeImplCreator : public Creator<::mmdeploy::FormatShapeImpl> {
- public:
-  const char* GetName() const override { return "cpu"; }
-  int GetVersion() const override { return 1; }
-  ReturnType Create(const Value& args) override { return make_unique<FormatShapeImpl>(args); }
-};
+MMDEPLOY_REGISTER_TRANSFORM_IMPL(::mmdeploy::FormatShapeImpl, (cpu, 0), FormatShapeImpl);
 
 }  // namespace cpu
 }  // namespace mmdeploy
-
-using ::mmdeploy::FormatShapeImpl;
-using ::mmdeploy::cpu::FormatShapeImplCreator;
-REGISTER_MODULE(FormatShapeImpl, FormatShapeImplCreator);
diff --git a/csrc/mmdeploy/codebase/mmaction/cuda/format_shape_impl.cpp b/csrc/mmdeploy/codebase/mmaction/cuda/format_shape_impl.cpp
@@ -114,16 +114,7 @@ class FormatShapeImpl : public ::mmdeploy::FormatShapeImpl {
   }
 };
 
-class FormatShapeImplCreator : public Creator<::mmdeploy::FormatShapeImpl> {
- public:
-  const char* GetName() const override { return "cuda"; }
-  int GetVersion() const override { return 1; }
-  ReturnType Create(const Value& args) override { return make_unique<FormatShapeImpl>(args); }
-};
+MMDEPLOY_REGISTER_TRANSFORM_IMPL(::mmdeploy::FormatShapeImpl, (cuda, 0), FormatShapeImpl);
 
 }  // namespace cuda
 }  // namespace mmdeploy
-
-using ::mmdeploy::FormatShapeImpl;
-using ::mmdeploy::cuda::FormatShapeImplCreator;
-REGISTER_MODULE(FormatShapeImpl, FormatShapeImplCreator);
diff --git a/csrc/mmdeploy/codebase/mmaction/format_shape.cpp b/csrc/mmdeploy/codebase/mmaction/format_shape.cpp
@@ -58,7 +58,7 @@ Result<Value> FormatShapeImpl::Process(const Value& input) {
 class FormatShape : public Transform {
  public:
   explicit FormatShape(const Value& args, int version = 0) : Transform(args) {
-    auto impl_creator = Registry<FormatShapeImpl>::Get().GetCreator(specified_platform_, version);
+    auto impl_creator = gRegistry<FormatShapeImpl>().Get(specified_platform_, version);
     if (nullptr == impl_creator) {
       MMDEPLOY_ERROR("'FormatShape' is not supported on '{}' platform", specified_platform_);
       throw std::domain_error("'FormatShape' is not supported on specified platform");
@@ -73,17 +73,10 @@ class FormatShape : public Transform {
   std::unique_ptr<FormatShapeImpl> impl_;
 };
 
-class FormatShapeCreator : public Creator<Transform> {
- public:
-  const char* GetName(void) const override { return "FormatShape"; }
-  int GetVersion(void) const override { return version_; }
-  ReturnType Create(const Value& args) override { return make_unique<FormatShape>(args, version_); }
-
- private:
-  int version_{1};
-};
+MMDEPLOY_REGISTER_FACTORY_FUNC(Transform, (FormatShape, 0), [](const Value& config) {
+  return std::make_unique<FormatShape>(config, 0);
+});
 
-REGISTER_MODULE(Transform, FormatShapeCreator);
 MMDEPLOY_DEFINE_REGISTRY(FormatShapeImpl);
 
 }  // namespace mmdeploy
diff --git a/csrc/mmdeploy/codebase/mmaction/format_shape.h b/csrc/mmdeploy/codebase/mmaction/format_shape.h
@@ -30,7 +30,7 @@ class FormatShapeImpl : public TransformImpl {
   ArgType arg_;
 };
 
-MMDEPLOY_DECLARE_REGISTRY(FormatShapeImpl);
+MMDEPLOY_DECLARE_REGISTRY(FormatShapeImpl, std::unique_ptr<FormatShapeImpl>(const Value& config));
 
 }  // namespace mmdeploy
 

diff --git a/csrc/mmdeploy/codebase/mmaction/mmaction.cpp b/csrc/mmdeploy/codebase/mmaction/mmaction.cpp
@@ -2,12 +2,8 @@
 
 #include "mmdeploy/codebase/mmaction/mmaction.h"
 
-namespace mmdeploy {
-namespace mmaction {
+namespace mmdeploy::mmaction {
 
-REGISTER_CODEBASE(MMAction);
+MMDEPLOY_REGISTER_CODEBASE(MMAction);
 
-}
-
-MMDEPLOY_DEFINE_REGISTRY(mmaction::MMAction);
-}  // namespace mmdeploy
+}  // namespace mmdeploy::mmaction