mirror of
https://github.com/open-mmlab/mmdeploy.git
synced 2025-01-14 08:09:43 +08:00
36124f6205
* 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>
57 lines
1.9 KiB
C++
57 lines
1.9 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "codebase/mmcls/mmcls.h"
|
|
#include "core/tensor.h"
|
|
#include "core/utils/formatter.h"
|
|
#include "experimental/module_adapter.h"
|
|
|
|
using std::vector;
|
|
|
|
namespace mmdeploy::mmcls {
|
|
|
|
class SoftmaxPost : public MMClsPostprocess {
|
|
public:
|
|
explicit SoftmaxPost(const Value& cfg) : MMClsPostprocess(cfg) {}
|
|
|
|
Result<Value> operator()(const Value& data, const Value& infer_res) {
|
|
DEBUG("data: {}, infer_res: {}", data, infer_res);
|
|
auto output_tensor = infer_res["cls"].get<Tensor>();
|
|
assert(output_tensor.shape().size() >= 2);
|
|
auto batch_size = (int)output_tensor.shape()[0];
|
|
auto class_num = (int)output_tensor.shape()[1];
|
|
|
|
if (output_tensor.device().is_host()) {
|
|
vector<float> scores(output_tensor.data<float>(),
|
|
output_tensor.data<float>() + output_tensor.size());
|
|
OUTCOME_TRY(stream().Wait());
|
|
return GetLabels(data, scores, batch_size, class_num);
|
|
} else {
|
|
vector<float> scores(output_tensor.size());
|
|
OUTCOME_TRY(output_tensor.CopyTo(scores.data(), stream()));
|
|
OUTCOME_TRY(stream().Wait());
|
|
return GetLabels(data, scores, batch_size, class_num);
|
|
}
|
|
}
|
|
|
|
private:
|
|
static Value GetLabels(const Value& data, const vector<float>& scores, int batch_size,
|
|
int class_num) {
|
|
ClassifyOutput output;
|
|
auto score_ptr = scores.data();
|
|
for (int i = 0; i < batch_size; ++i, score_ptr += class_num) {
|
|
auto max_score_ptr = std::max_element(score_ptr, score_ptr + class_num);
|
|
ClassifyOutput::Label label{int(max_score_ptr - score_ptr), *max_score_ptr};
|
|
DEBUG("label_id: {}, score: {}", label.label_id, label.score);
|
|
output.labels.push_back(label);
|
|
}
|
|
return to_value(std::move(output));
|
|
}
|
|
|
|
private:
|
|
float thres_{0.f};
|
|
};
|
|
|
|
REGISTER_CODEBASE_MODULE(MMClsPostprocess, SoftmaxPost);
|
|
|
|
} // namespace mmdeploy::mmcls
|