nms.cpp 1.2 KB
Newer Older
1
/**
M
Megvii Engine Team 已提交
2 3
 * \file imperative/src/impl/ops/nms.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
4
 *
M
Megvii Engine Team 已提交
5
 * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
6
 *
M
Megvii Engine Team 已提交
7 8 9
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10
 */
M
Megvii Engine Team 已提交
11

12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45
#include "../op_trait.h"

#include "megbrain/imperative/ops/nms.h"
#include "megbrain/opr/standalone/nms_opr.h"

namespace mgb {
namespace imperative {

using NMSKeepOpr = opr::standalone::NMSKeep;

namespace {
cg::OperatorNodeBase* apply_on_var_node(
        const OpDef& def,
        const VarNodeArray& inputs) {
    auto&& nms_keep = def.cast_final_safe<NMSKeep>();

    NMSKeepOpr::Param param;
    param.iou_thresh = nms_keep.iou_thresh;
    param.max_output = nms_keep.max_output;

    return NMSKeepOpr::make(inputs[0], param).node()->owner_opr();
}

OP_TRAIT_REG(NMSKeep, NMSKeep, NMSKeepOpr)
    .apply_on_var_node(apply_on_var_node)
    .fallback();
} // anonymous namespace

MGB_DYN_TYPE_OBJ_FINAL_IMPL(NMSKeep);

}  // namespace imperative
}  // namespace mgb

// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}