wehub-resource-sync
122 lines
5.0 KiB
C++
122 lines
5.0 KiB
C++
//
|
|
// OnnxNonMaxSuppression.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2020/04/01.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <flatbuffers/util.h>
|
|
#include "MNN_generated.h"
|
|
#include "OnnxExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class OnnxNonMaxSuppressionTransformer : public OnnxExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto op = expr->get();
|
|
MNN_ASSERT(op->type() == OpType_Extra);
|
|
|
|
VARP result(nullptr);
|
|
// NonMaxSuppression(boxes, scores, max_output_size, iou_threshold, score_threshold)
|
|
// boxes and scores are required, the last 3 parameters are optional.
|
|
// onnx boxes is 3D [num_batches, boxes_num, 4] with num_batches = 1, while tf boxes
|
|
// is 2D [boxes_num, 4].
|
|
// onnx scores is 3D [num_batches, num_classes, boxes_num] with num_batches = 1,
|
|
// while tf scores is 1D [boxes_num].
|
|
auto inputs = expr->inputs();
|
|
// optional input 3/4/5th
|
|
if (inputs.size() < 3 || inputs[2].get() == nullptr) {
|
|
MNN_ERROR("NonMaxSuppression's max_output_boxes_per_class must be provided (can't optional)\n");
|
|
return nullptr;
|
|
}
|
|
auto zero = _Scalar<float>(0);
|
|
for (int i = 3; i < inputs.size(); ++i) {
|
|
if (inputs[i].get() == nullptr) {
|
|
inputs[i] = zero;
|
|
}
|
|
}
|
|
|
|
auto input0Info = inputs[0]->getInfo();
|
|
auto input1Info = inputs[1]->getInfo();
|
|
bool oldSupport = (input0Info != nullptr && input1Info != nullptr);
|
|
if (oldSupport) {
|
|
for (auto dim : input0Info->dim) {
|
|
if (dim <= 0) {
|
|
oldSupport = false;
|
|
break;
|
|
}
|
|
}
|
|
for (auto dim : input1Info->dim) {
|
|
if (dim <= 0) {
|
|
oldSupport = false;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (!oldSupport) {
|
|
MNN_ERROR("Shape of NonMaxSupression's input is unknown. Please confirm version of MNN engine is new enough and use V3 Module API to run it correctly\n");
|
|
std::unique_ptr<OpT> nms(new OpT);
|
|
nms->type = OpType_NonMaxSuppressionV2;
|
|
nms->main.type = OpParameter_NonMaxSuppressionV2;
|
|
nms->main.value = new NonMaxSuppressionV2T;
|
|
auto result = Expr::create(nms.get(), inputs);
|
|
Variable::create(result)->setName(expr->outputName(0));
|
|
return result;
|
|
}
|
|
MNN_ASSERT(inputs[0]->getInfo()->dim.size() == 3);
|
|
MNN_ASSERT(inputs[1]->getInfo()->dim.size() == 3);
|
|
|
|
for (int batch = 0; batch < inputs[0]->getInfo()->dim[0]; ++batch) {
|
|
VARP boxes = _Gather(inputs[0], _Scalar<int>(batch)); // [boxes_num, 4]
|
|
VARP scores = _Gather(inputs[1], _Scalar<int>(batch)); // [num_classes, boxes_num]
|
|
|
|
int num_classes = scores->getInfo()->dim[0];
|
|
for (int cls = 0; cls < num_classes; ++cls) {
|
|
VARP scores_per_class = _Gather(scores, _Scalar<int>(cls)); // [boxes_num]
|
|
|
|
std::unique_ptr<MNN::OpT> nonMaxSuppressionOp(new OpT);
|
|
std::string name = op->name()->str() + "/" + flatbuffers::NumToString(cls);
|
|
nonMaxSuppressionOp->name = name;
|
|
nonMaxSuppressionOp->type = OpType_NonMaxSuppressionV2;
|
|
nonMaxSuppressionOp->main.type = OpParameter_NonMaxSuppressionV2;
|
|
nonMaxSuppressionOp->main.value = nullptr;
|
|
|
|
std::vector<VARP> newInputs{boxes, scores_per_class};
|
|
for (int i = 2; i < inputs.size(); ++i) {
|
|
newInputs.push_back(inputs[i]);
|
|
}
|
|
auto nonMaxSupp = Expr::create(nonMaxSuppressionOp.get(), newInputs, 1 /*output size*/);
|
|
nonMaxSupp->setName(expr->name() + "/" + flatbuffers::NumToString(cls));
|
|
|
|
// Tensorflow's output is [num_selected_boxes], while onnx requires
|
|
// [num_selected_boxes, 3], and the meaning of last dim is
|
|
// [batch_index, class_index, box_index].
|
|
VARP output = _Unsqueeze(Variable::create(nonMaxSupp), {1}); // [num_selected_boxes, 1]
|
|
|
|
auto shape = _Shape(output, true);
|
|
output = _Concat({_Fill(shape, _Scalar<int>(batch)), _Fill(shape, _Scalar<int>(cls)), output}, 1);
|
|
if (result.get() != nullptr) {
|
|
result = _Concat({result, output}, 0);
|
|
} else {
|
|
result = output;
|
|
}
|
|
}
|
|
}
|
|
result->setName(expr->outputName(0));
|
|
return result->expr().first;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
OnnxExtraManager::get()->insert("NonMaxSuppression",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxNonMaxSuppressionTransformer));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|