122 lines
3.7 KiB
C++
122 lines
3.7 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "restorer.h"
|
|
|
|
#include "codebase/mmedit/mmedit.h"
|
|
#include "core/device.h"
|
|
#include "core/graph.h"
|
|
#include "core/mat.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", "det"},
|
|
{"type", "Inference"},
|
|
{"params", {{"model", "TBD"}}},
|
|
{"input", {"img"}},
|
|
{"output", {"out"}}
|
|
}
|
|
}
|
|
},
|
|
{"input", {"img"}},
|
|
{"output", {"out"}}
|
|
}
|
|
}
|
|
};
|
|
// clang-format on
|
|
return v;
|
|
}
|
|
|
|
template <class ModelType>
|
|
int mmdeploy_restorer_create_impl(ModelType &&m, const char *device_name, int device_id,
|
|
mm_handle_t *handle) {
|
|
try {
|
|
auto config = config_template();
|
|
config["pipeline"]["tasks"][0]["params"]["model"] = std::forward<ModelType>(m);
|
|
|
|
auto restorer = std::make_unique<Handle>(device_name, device_id, std::move(config));
|
|
|
|
*handle = restorer.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_restorer_create(mm_model_t model, const char *device_name, int device_id,
|
|
mm_handle_t *handle) {
|
|
return mmdeploy_restorer_create_impl(*static_cast<Model *>(model), device_name, device_id,
|
|
handle);
|
|
}
|
|
|
|
int mmdeploy_restorer_create_by_path(const char *model_path, const char *device_name, int device_id,
|
|
mm_handle_t *handle) {
|
|
return mmdeploy_restorer_create_impl(model_path, device_name, device_id, handle);
|
|
}
|
|
|
|
int mmdeploy_restorer_apply(mm_handle_t handle, const mm_mat_t *images, int count,
|
|
mm_mat_t **results) {
|
|
if (handle == nullptr || images == nullptr || count == 0 || results == nullptr) {
|
|
return MM_E_INVALID_ARG;
|
|
}
|
|
try {
|
|
auto restorer = static_cast<Handle *>(handle);
|
|
Value input{Value::kArray};
|
|
for (int i = 0; i < count; ++i) {
|
|
Mat _mat{images[i].height, images[i].width, PixelFormat(images[i].format),
|
|
DataType(images[i].type), images[i].data, Device{"cpu"}};
|
|
input.front().push_back({{"ori_img", _mat}});
|
|
}
|
|
auto output = restorer->Run(std::move(input)).value().front();
|
|
auto restorer_output = from_value<std::vector<mmedit::RestorerOutput>>(output);
|
|
|
|
auto deleter = [&](mm_mat_t *p) { mmdeploy_restorer_release_result(p, count); };
|
|
|
|
std::unique_ptr<mm_mat_t[], decltype(deleter)> _results(new mm_mat_t[count]{}, deleter);
|
|
|
|
for (int i = 0; i < count; ++i) {
|
|
auto upscale = restorer_output[i];
|
|
auto &res = _results[i];
|
|
res.data = new uint8_t[upscale.byte_size()];
|
|
memcpy(res.data, upscale.data<uint8_t>(), upscale.byte_size());
|
|
res.format = (mm_pixel_format_t)upscale.pixel_format();
|
|
res.height = upscale.height();
|
|
res.width = upscale.width();
|
|
res.type = (mm_data_type_t)upscale.type();
|
|
}
|
|
*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_restorer_release_result(mm_mat_t *results, int count) {
|
|
for (int i = 0; i < count; ++i) {
|
|
delete[] results[i].data;
|
|
}
|
|
delete[] results;
|
|
}
|
|
|
|
void mmdeploy_restorer_destroy(mm_handle_t handle) { delete static_cast<Handle *>(handle); }
|