Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!7681 nms registry 

Merge pull request !7681 from zhaozhenlong/lite/issue/nms_registry
This commit is contained in:
mindspore-ci-bot 2020-10-23 17:25:31 +08:00 committed by Gitee
parent cf06b22f3d a43fe26a01
commit d38370cb39
2 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
mindspore/lite/src/ops/non_max_suppression.cc
View File

@ -16,6 +16,10 @@
#include "src/ops/non_max_suppression.h"
#ifndef PRIMITIVE_WRITEABLE
#include "src/ops/ops_register.h"
#endif
namespace mindspore {
namespace lite {
#ifdef PRIMITIVE_WRITEABLE
@ -44,6 +48,13 @@ int NonMaxSuppression::UnPackToFlatBuilder(const schema::Primitive *primitive, f
int NonMaxSuppression::GetCenterPointBox() const {
return this->primitive_->value_as_NonMaxSuppression()->centerPointBox();
}
PrimitiveC *NonMaxSuppressionCreator(const schema::Primitive *primitive) {
return PrimitiveC::NewPrimitiveC<NonMaxSuppression>(primitive);
}
Registry NonMaxSuppressionRegistry(schema::PrimitiveType_NonMaxSuppression, NonMaxSuppressionCreator);
#endif
int NonMaxSuppression::InferShape(std::vector<Tensor *> inputs_, std::vector<Tensor *> outputs_) {
auto input = inputs_.front();

3
mindspore/lite/src/ops/populate/non_max_suppression_populate.cc
View File

@ -35,7 +35,8 @@ OpParameter *PopulateNonMaxSuppressionParameter(const mindspore::lite::Primitive
param->center_point_box_ = prim->GetCenterPointBox();
return reinterpret_cast<OpParameter *>(param);
}
Registry NonMaxSuppressionParameterRegistry(schema::PrimitiveType_OneHot, PopulateNonMaxSuppressionParameter);
Registry NonMaxSuppressionParameterRegistry(schema::PrimitiveType_NonMaxSuppression,
PopulateNonMaxSuppressionParameter);
} // namespace lite
} // namespace mindspore