mmdeploy/csrc/apis/python/pose_detector.cpp
lzhangzz 46bfe0ac87
[Feature] New pipeline & executor for SDK (#497)
* executor prototype

* add split/when_all

* fix GCC build

* WIP let_value

* fix let_value

* WIP ensure_started

* ensure_started & start_detached

* fix let_value + when_all combo on MSVC 142

* fix static thread pool

* generic just, then, let_value, sync_wait

* minor

* generic split and when_all

* fully generic sender adapters

* when_all: workaround for GCC7

* support legacy spdlog

* fix memleak

* bulk

* static detector

* fix bulk & first pipeline

* bulk for static thread pools

* fix on MSVC

* WIP async batch submission

* WIP collation

* async batch

* fix detector

* fix async detector

* fix

* fix

* debug

* fix cuda allocator

* WIP type erased executor

* better type erasure

* simplify C API impl

* Expand & type erase TC

* deduction guide for type erased senders

* fix GCC build

* when_all for arrays of Value senders

* WIP pipeline v2

* WIP pipeline parser

* WIP timed batch operation

* add registry

* experiment

* fix pipeline

* naming

* fix mem-leak

* fix deferred batch operation

* WIP

* WIP configurable scheduler

* WIP configurable scheduler

* add comment

* parse scheduler config

* force link schedulers

* WIP pipeable sender

* WIP CPO

* ADL isolation and dismantle headers

* type erase single thread context

* fix MSVC build

* CPO

* replace decay_t with remove_cvref_t

* structure adjustment

* structure adjustment

* apply CPOs & C API rework

* refine C API

* detector async C API

* adjust detector async C API

* # Conflicts:
#	csrc/apis/c/detector.cpp

* fix when_all for type erased senders

* support void return for Then

* async detector

* fix some CPOs

* minor

* WIP rework capture mechanism for type erased types

* minor fix

* fix MSVC build

* move expand.h to execution

* make `Expand` pipeable

* fix type erased

* un-templatize `_TypeErasedOperation`

* re-work C API

* remove async_detector C API

* fix pipeline

* add flatten & unflatten

* fix flatten & unflatten

* add aync OCR demo

* config executor for nodes & better executor API

* working async OCR example

* minor

* dynamic batch via scheduler

* dynamic batch on `Value`

* fix MSVC build

* type erase dynamic batch scheduler

* sender as Python Awaitable

* naming

* naming

* add docs

* minor

* merge tmp branch

* unify C APIs

* fix ocr

* unify APIs

* fix typo

* update async OCR demo

* add v3 API text recognizer

* fix v3 API

* fix lint

* add license info & reformat

* add demo async_ocr_v2

* revert files

* revert files

* resolve link issues

* fix scheduler linkage for shared libs

* fix license header

* add docs for `mmdeploy_executor_split`

* add missing `mmdeploy_executor_transfer_just` and `mmdeploy_executor_execute`

* make `TimedSingleThreadContext` header only

* fix lint

* simplify type-erased sender
2022-06-01 14:10:43 +08:00

121 lines
3.6 KiB
C++

// Copyright (c) OpenMMLab. All rights reserved.
#include "pose_detector.h"
#include <array>
#include "common.h"
#include "core/logger.h"
namespace mmdeploy {
using Rect = std::array<float, 4>;
class PyPoseDedector {
public:
PyPoseDedector(const char *model_path, const char *device_name, int device_id) {
MMDEPLOY_INFO("{}, {}, {}", model_path, device_name, device_id);
auto status =
mmdeploy_pose_detector_create_by_path(model_path, device_name, device_id, &handle_);
if (status != MM_SUCCESS) {
throw std::runtime_error("failed to create pose_detector");
}
}
py::list Apply(const std::vector<PyImage> &imgs, const std::vector<std::vector<Rect>> &vboxes) {
if (imgs.size() == 0 && vboxes.size() == 0) {
return py::list{};
}
if (vboxes.size() != 0 && vboxes.size() != imgs.size()) {
std::string error =
fmt::format("imgs length not equal with vboxes [{} vs {}]", imgs.size(), vboxes.size());
throw std::invalid_argument(error);
}
std::vector<mm_mat_t> mats;
std::vector<mm_rect_t> boxes;
std::vector<int> bbox_count;
mats.reserve(imgs.size());
for (const auto &img : imgs) {
auto mat = GetMat(img);
mats.push_back(mat);
}
for (auto _boxes : vboxes) {
for (auto _box : _boxes) {
mm_rect_t box = {_box[0], _box[1], _box[2], _box[3]};
boxes.push_back(box);
}
bbox_count.push_back(_boxes.size());
}
// full image
if (vboxes.size() == 0) {
for (int i = 0; i < mats.size(); i++) {
mm_rect_t box = {0.f, 0.f, mats[i].width - 1, mats[i].height - 1};
boxes.push_back(box);
bbox_count.push_back(1);
}
}
mm_pose_detect_t *detection{};
auto status = mmdeploy_pose_detector_apply_bbox(handle_, mats.data(), (int)mats.size(),
boxes.data(), bbox_count.data(), &detection);
if (status != MM_SUCCESS) {
throw std::runtime_error("failed to apply pose_detector, code: " + std::to_string(status));
}
auto output = py::list{};
auto result = detection;
for (int i = 0; i < mats.size(); i++) {
if (bbox_count[i] == 0) {
output.append(py::none());
continue;
}
int n_point = result->length;
auto pred = py::array_t<float>({bbox_count[i], n_point, 3});
auto dst = pred.mutable_data();
for (int j = 0; j < bbox_count[i]; j++) {
for (int k = 0; k < n_point; k++) {
dst[0] = result->point[k].x;
dst[1] = result->point[k].y;
dst[2] = result->score[k];
dst += 3;
}
result++;
}
output.append(std::move(pred));
}
int total = std::accumulate(bbox_count.begin(), bbox_count.end(), 0);
mmdeploy_pose_detector_release_result(detection, total);
return output;
}
~PyPoseDedector() {
mmdeploy_pose_detector_destroy(handle_);
handle_ = {};
}
private:
mm_handle_t handle_{};
};
static void register_python_pose_detector(py::module &m) {
py::class_<PyPoseDedector>(m, "PoseDetector")
.def(py::init([](const char *model_path, const char *device_name, int device_id) {
return std::make_unique<PyPoseDedector>(model_path, device_name, device_id);
}))
.def("__call__", &PyPoseDedector::Apply, py::arg("imgs"),
py::arg("vboxes") = std::vector<std::vector<Rect>>());
}
class PythonPoseDetectorRegisterer {
public:
PythonPoseDetectorRegisterer() {
gPythonBindings().emplace("pose_detector", register_python_pose_detector);
}
};
static PythonPoseDetectorRegisterer python_pose_detector_registerer;
} // namespace mmdeploy