mirror of
https://github.com/open-mmlab/mmdeploy.git
synced 2025-01-14 08:09:43 +08:00
* check in cmake * move backend_ops to csrc/backend_ops * check in preprocess, model, some codebase and their c-apis * check in CMakeLists.txt * check in parts of test_csrc * commit everything else * add readme * update core's BUILD_INTERFACE directory * skip codespell on third_party * update trt_net and ort_net's CMakeLists * ignore clion's build directory * check in pybind11 * add onnx.proto. Remove MMDeploy's dependency on ncnn's source code * export MMDeployTargets only when MMDEPLOY_BUILD_SDK is ON * remove useless message * target include directory is wrong * change target name from mmdeploy_ppl_net to mmdeploy_pplnn_net * skip install directory * update project's cmake * remove useless code * set CMAKE_BUILD_TYPE to Release by force if it isn't set by user * update custom ops CMakeLists * pass object target's source lists * fix lint end-of-file * fix lint: trailing whitespace * fix codespell hook * remove bicubic_interpolate to csrc/backend_ops/ * set MMDEPLOY_BUILD_SDK OFF * change custom ops build command * add spdlog installation command * update docs on how to checkout pybind11 * move bicubic_interpolate to backend_ops/tensorrt directory * remove useless code * correct cmake * fix typo * fix typo * fix install directory * correct sdk's readme * set cub dir when cuda version < 11.0 * change directory where clang-format will apply to * fix build command * add .clang-format * change clang-format style from google to file * reformat csrc/backend_ops * format sdk's code * turn off clang-format for some files * add -Xcompiler=-fno-gnu-unique * fix trt topk initialize * check in config for sdk demo * update cmake script and csrc's readme * correct config's path * add cuda include directory, otherwise compile failed in case of tensorrt8.2 * clang-format onnx2ncnn.cpp Co-authored-by: zhangli <lzhang329@gmail.com> Co-authored-by: grimoire <yaoqian@sensetime.com>
197 lines
6.1 KiB
C++
197 lines
6.1 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "text_recognizer.h"
|
|
|
|
#include <numeric>
|
|
|
|
#include "archive/value_archive.h"
|
|
#include "codebase/mmocr/mmocr.h"
|
|
#include "core/device.h"
|
|
#include "core/mat.h"
|
|
#include "core/operator.h"
|
|
#include "core/status_code.h"
|
|
#include "core/utils/formatter.h"
|
|
#include "handle.h"
|
|
|
|
using namespace mmdeploy;
|
|
|
|
namespace {
|
|
|
|
const Value &config_template() {
|
|
// clang-format off
|
|
static Value v {
|
|
{
|
|
"pipeline", {
|
|
{
|
|
"tasks", {
|
|
{
|
|
{"name", "warp"},
|
|
{"type", "Task"},
|
|
{"module", "WarpBoxes"},
|
|
{"input", {"img", "dets"}},
|
|
{"output", {"patches"}}
|
|
},
|
|
{
|
|
{"name", "flatten"},
|
|
{"type", "Flatten"},
|
|
{"input", {"patches"}},
|
|
{"output", {"patch_flat", "patch_index"}},
|
|
},
|
|
{
|
|
{"name", "recog"},
|
|
{"type", "Inference"},
|
|
{"params", {{"model", "TBD"},{"batch_size", 1}}},
|
|
{"input", {"patch_flat"}},
|
|
{"output", {"texts"}}
|
|
},
|
|
{
|
|
{"name", "unflatten"},
|
|
{"type", "Unflatten"},
|
|
{"input", {"texts", "patch_index"}},
|
|
{"output", {"text_unflat"}},
|
|
}
|
|
}
|
|
},
|
|
{"input", {"img", "dets"}},
|
|
{"output", {"text_unflat"}}
|
|
}
|
|
}
|
|
};
|
|
// clang-format on
|
|
return v;
|
|
}
|
|
|
|
template <class ModelType>
|
|
int mmdeploy_text_recognizer_create_impl(ModelType &&m, const char *device_name, int device_id,
|
|
mm_handle_t *handle) {
|
|
try {
|
|
auto value = config_template();
|
|
value["pipeline"]["tasks"][2]["params"]["model"] = std::forward<ModelType>(m);
|
|
|
|
auto recognizer = std::make_unique<Handle>(device_name, device_id, std::move(value));
|
|
|
|
*handle = recognizer.release();
|
|
return MM_SUCCESS;
|
|
|
|
} catch (const std::exception &e) {
|
|
ERROR("exception caught: {}", e.what());
|
|
} catch (...) {
|
|
ERROR("unknown exception caught");
|
|
}
|
|
return MM_E_FAIL;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int mmdeploy_text_recognizer_create(mm_model_t model, const char *device_name, int device_id,
|
|
mm_handle_t *handle) {
|
|
return mmdeploy_text_recognizer_create_impl(*static_cast<Model *>(model), device_name, device_id,
|
|
handle);
|
|
}
|
|
|
|
int mmdeploy_text_recognizer_create_by_path(const char *model_path, const char *device_name,
|
|
int device_id, mm_handle_t *handle) {
|
|
return mmdeploy_text_recognizer_create_impl(model_path, device_name, device_id, handle);
|
|
}
|
|
|
|
int mmdeploy_text_recognizer_apply(mm_handle_t handle, const mm_mat_t *images, int count,
|
|
mm_text_recognize_t **results) {
|
|
return mmdeploy_text_recognizer_apply_bbox(handle, images, count, nullptr, nullptr, results);
|
|
}
|
|
|
|
int mmdeploy_text_recognizer_apply_bbox(mm_handle_t handle, const mm_mat_t *images, int image_count,
|
|
const mm_text_detect_t *bboxes, const int *bbox_count,
|
|
mm_text_recognize_t **results) {
|
|
if (handle == nullptr || images == nullptr || image_count == 0 || results == nullptr) {
|
|
return MM_E_INVALID_ARG;
|
|
}
|
|
|
|
try {
|
|
auto recognizer = static_cast<Handle *>(handle);
|
|
Value input{Value::kArray, Value::kArray};
|
|
auto _bboxes = bboxes;
|
|
auto result_count = 0;
|
|
for (int i = 0; i < image_count; ++i) {
|
|
mmdeploy::Mat _mat{images[i].height, images[i].width, PixelFormat(images[i].format),
|
|
DataType(images->type), images[i].data, Device{"cpu"}};
|
|
input[0].push_back({{"ori_img", _mat}});
|
|
|
|
if (bboxes && bbox_count) {
|
|
Value boxes(Value::kArray);
|
|
for (int j = 0; j < bbox_count[i]; ++j) {
|
|
Value box;
|
|
for (const auto &p : _bboxes[j].bbox) {
|
|
box.push_back(p.x);
|
|
box.push_back(p.y);
|
|
}
|
|
boxes.push_back(std::move(box));
|
|
}
|
|
_bboxes += bbox_count[i];
|
|
result_count += bbox_count[i];
|
|
input[1].push_back({{"boxes", boxes}});
|
|
} else {
|
|
input[1].push_back(Value::kNull);
|
|
result_count += 1;
|
|
}
|
|
}
|
|
|
|
auto output = recognizer->Run(std::move(input)).value().front();
|
|
|
|
auto recognizer_outputs =
|
|
from_value<std::vector<std::vector<mmocr::TextRecognizerOutput>>>(output);
|
|
|
|
std::vector<int> counts;
|
|
if (bboxes && bbox_count) {
|
|
counts = std::vector<int>(bbox_count, bbox_count + image_count);
|
|
} else {
|
|
counts.resize(image_count, 1);
|
|
}
|
|
std::vector<int> offsets{0};
|
|
std::partial_sum(begin(counts), end(counts), back_inserter(offsets));
|
|
|
|
auto deleter = [&](mm_text_recognize_t *p) {
|
|
mmdeploy_text_recognizer_release_result(p, offsets.back());
|
|
};
|
|
|
|
std::unique_ptr<mm_text_recognize_t[], decltype(deleter)> _results(
|
|
new mm_text_recognize_t[result_count]{}, deleter);
|
|
|
|
for (int i = 0; i < image_count; ++i) {
|
|
auto &recog_output = recognizer_outputs[i];
|
|
for (int j = 0; j < recog_output.size(); ++j) {
|
|
auto &res = _results[offsets[i] + j];
|
|
|
|
auto &box_result = recog_output[j];
|
|
|
|
auto &score = box_result.score;
|
|
res.length = static_cast<int>(score.size());
|
|
|
|
res.score = new float[score.size()];
|
|
std::copy_n(score.data(), score.size(), res.score);
|
|
|
|
auto text = box_result.text;
|
|
res.text = new char[text.length() + 1];
|
|
std::copy_n(text.data(), text.length() + 1, res.text);
|
|
}
|
|
}
|
|
*results = _results.release();
|
|
return MM_SUCCESS;
|
|
|
|
} catch (const std::exception &e) {
|
|
ERROR("exception caught: {}", e.what());
|
|
} catch (...) {
|
|
ERROR("unknown exception caught");
|
|
}
|
|
return MM_E_FAIL;
|
|
}
|
|
|
|
void mmdeploy_text_recognizer_release_result(mm_text_recognize_t *results, int count) {
|
|
for (int i = 0; i < count; ++i) {
|
|
delete[] results[i].score;
|
|
delete[] results[i].text;
|
|
}
|
|
delete[] results;
|
|
}
|
|
|
|
void mmdeploy_text_recognizer_destroy(mm_handle_t handle) { delete static_cast<Handle *>(handle); }
|