From 412c2f35a9ef5cc766c3a7e41bee75eb6c407221 Mon Sep 17 00:00:00 2001 From: Hector Li Date: Thu, 29 Nov 2018 11:02:20 -0800 Subject: [PATCH 1/8] Add NonMaxSupression op to contribution ops --- onnxruntime/contrib_ops/contrib_ops.cc | 35 +++ .../contrib_ops/cpu/non_max_suppression.cc | 142 ++++++++++++ .../contrib_ops/cpu/non_max_suppression.h | 34 +++ .../contrib_ops/non_max_suppression_test.cc | 205 ++++++++++++++++++ 4 files changed, 416 insertions(+) create mode 100644 onnxruntime/contrib_ops/cpu/non_max_suppression.cc create mode 100644 onnxruntime/contrib_ops/cpu/non_max_suppression.h create mode 100644 onnxruntime/test/contrib_ops/non_max_suppression_test.cc diff --git a/onnxruntime/contrib_ops/contrib_ops.cc b/onnxruntime/contrib_ops/contrib_ops.cc index 59191395ac53e..8a295df3d8b31 100644 --- a/onnxruntime/contrib_ops/contrib_ops.cc +++ b/onnxruntime/contrib_ops/contrib_ops.cc @@ -357,6 +357,39 @@ with the exception that numpy default keepdims to False instead of True.)DOC") "keepdims", "Keep the reduced dimension or not, default 1 mean keep reduced dimension.", AttributeProto::INT); + + ONNX_CONTRIB_OPERATOR_SCHEMA(NonMaxSuppression) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc(R"DOC( +Pruning away boxes that have high intersection-over-union (IOU) overlap with previously selected boxes. +Bounding boxes with score less than score_threshold are removed. Bounding boxes are supplied as [y1, x1, y2, x2], +where (y1, x1) and (y2, x2) are the coordinates of any diagonal pair of box corners and the coordinates can be provided +as normalized (i.e., lying in the interval [0, 1]) or absolute. +Note that this algorithm is agnostic to where the origin is in the coordinate system and more generally is invariant to +orthogonal transformations and translations of the coordinate system; +thus translating or reflections of the coordinate system result in the same boxes being selected by the algorithm. +The output of this operation is a set of integers indexing into the input collection of bounding boxes representing the selected boxes. +The bounding box coordinates corresponding to the selected indices can then be obtained using the gather operation.)DOC") + .Input(0, "boxes", "An input tensor. 2D tensor with shape [num_boxes, 4]", "T1") + .Input(1, "scores", "An input tensor. 1D tensor with shape [num_boxes]", "T1") + .Output(0, "selected_indices", "selected indices from the boxes tensor.", "T2") + .TypeConstraint("T1", {"tensor(float)"}, "Constrain input type to float tensor.") + .TypeConstraint("T2", + {"tensor(int32)"}, + "Constrain output data type to 32-bit integer tensor.") + .Attr( + "max_output_size", + "Integer representing the maximum number of boxes to be selected by non max suppression.", + AttributeProto::INT) + .Attr( + "iou_threshold", + "Float representing the threshold for deciding whether boxes overlap too much with respect to IOU.", + AttributeProto::FLOAT) + .Attr( + "score_threshold", + "Float tensor representing the threshold for deciding when to remove boxes based on score.", + AttributeProto::FLOAT); } class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SampleOp); @@ -366,6 +399,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, DequantizeLinear); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int8_t, DequantizeLinear); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, QuantizeLinear); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, NonMaxSuppression); void RegisterContribKernels(std::function fn) { fn(BuildKernel()); @@ -378,6 +412,7 @@ void RegisterContribKernels(std::function fn) { fn(BuildKernel()); fn(BuildKernel()); fn(BuildKernel()); + fn(BuildKernel()); } } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/non_max_suppression.cc b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc new file mode 100644 index 0000000000000..b77d8880a42a5 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc @@ -0,0 +1,142 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cpu/non_max_suppression.h" +#include + +namespace onnxruntime { +namespace contrib { + +ONNX_CPU_OPERATOR_TYPED_MS_KERNEL( + NonMaxSuppression, + 1, + float, + KernelDefBuilder() + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()), + NonMaxSuppression); + +template +void NonMaxSuppression::MaxMin(const T& lhs, const T& rhs, T& min, T& max) const { + if (lhs >= rhs) { + min = rhs; + max = lhs; + } else { + min = lhs; + max = rhs; + } +} + +template +bool NonMaxSuppression::SuppressByIOU(const T* boxes_data, int32_t box_index1, int32_t box_index2) const { + T x1_min, y1_min, x1_max, y1_max, x2_min, y2_min, x2_max, y2_max; + // boxes data [y1, x1, y2, x2], + MaxMin(boxes_data[4 * box_index1 + 1], boxes_data[4 * box_index1 + 3], x1_min, x1_max); + MaxMin(boxes_data[4 * box_index1 + 0], boxes_data[4 * box_index1 + 2], y1_min, y1_max); + MaxMin(boxes_data[4 * box_index2 + 1], boxes_data[4 * box_index2 + 3], x2_min, x2_max); + MaxMin(boxes_data[4 * box_index2 + 0], boxes_data[4 * box_index2 + 2], y2_min, y2_max); + + const T intersection_x_min = std::max(x1_min, x2_min); + const T intersection_y_min = std::max(y1_min, y2_min); + const T intersection_x_max = std::min(x1_max, x2_max); + const T intersection_y_max = std::min(y1_max, y2_max); + + const T intersection_area = std::max(intersection_x_max - intersection_x_min, static_cast(0.0)) * + std::max(intersection_y_max - intersection_y_min, static_cast(0.0)); + + if (intersection_area <= static_cast(0.0)) { + return false; + } + + const T area1 = (x1_max - x1_min) * (y1_max - y1_min); + const T area2 = (x2_max - x2_min) * (y2_max - y2_min); + const T union_area = area1 + area2 - intersection_area; + + if (area1 <= static_cast(0.0) || area2 <= static_cast(0.0) || union_area <= static_cast(0.0)) { + return false; + } + + const T intersection_over_union = intersection_area / union_area; + + return intersection_over_union > iou_threshold_; +} + +template +Status NonMaxSuppression::Compute(OpKernelContext* ctx) const { + const Tensor* boxes = ctx->Input(0); + ONNXRUNTIME_ENFORCE(boxes); + const Tensor* scores = ctx->Input(1); + ONNXRUNTIME_ENFORCE(scores); + + const TensorShape& boxes_shape = boxes->Shape(); + auto boxes_dims = boxes_shape.GetDims(); + ONNXRUNTIME_RETURN_IF_NOT(boxes_shape.NumDimensions() == 2, "boxes must be a 2D tensor."); + int64_t num_boxes = boxes_dims[0]; + ONNXRUNTIME_RETURN_IF_NOT(boxes_dims[1] == 4, "boxes shape must be a 2D tensor with shape [num_boxes, 4]."); + + const TensorShape& scores_shape = scores->Shape(); + ONNXRUNTIME_RETURN_IF_NOT(scores_shape.NumDimensions() == 1, "boxes must be a 1D tensor."); + ONNXRUNTIME_RETURN_IF_NOT(scores_shape.GetDims()[0] == num_boxes, "scores and boxes should have same num_boxes."); + + if (max_output_size_ <= 0 || boxes_dims[0] == 0) { + std::vector output_dims(1, 0); + TensorShape output_shape(output_dims); + auto output_tensor = ctx->Output(0, output_shape); + return Status::OK(); + } + + const T* boxes_data = boxes->Data(); + const T* scores_data = scores->Data(); + + struct ScoreIndexPair { + T score; + int32_t index; + }; + + auto LessCompare = [](const ScoreIndexPair& lhs, const ScoreIndexPair& rhs) { + return lhs.score < rhs.score; + }; + + // Filter by score_threshold_ + std::priority_queue, decltype(LessCompare)> sorted_scores_with_index(LessCompare); + for (int32_t i = 0; i < num_boxes; ++i) { + if (static_cast(scores_data[i]) > score_threshold_) { + sorted_scores_with_index.emplace(ScoreIndexPair({scores_data[i], i})); + } + } + + int num_of_selected = 0; + std::vector selected_index(max_output_size_, 0); + ScoreIndexPair next_top_score; + + // Get the next box with top score, filter by iou_threshold_ + while (num_of_selected < max_output_size_ && !sorted_scores_with_index.empty()) { + next_top_score = sorted_scores_with_index.top(); + sorted_scores_with_index.pop(); + + bool selected = true; + // Check with existing boxes, suppress if exceed the IOU (Intersection Over Union) threadhold + for (int i = num_of_selected - 1; i >= 0; --i) { + if (SuppressByIOU(boxes_data, selected_index[i], next_top_score.index)) { + selected = false; + break; + } + } + + if (selected) { + selected_index[num_of_selected] = next_top_score.index; + ++num_of_selected; + } + } + + std::vector output_dim(1, num_of_selected); + TensorShape output_shape(output_dim); + Tensor* selected_indices = ctx->Output(0, output_shape); + auto output_data = selected_indices->MutableData(); + memcpy(output_data, selected_index.data(), num_of_selected * sizeof(int32_t)); + + return Status::OK(); +} + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/non_max_suppression.h b/onnxruntime/contrib_ops/cpu/non_max_suppression.h new file mode 100644 index 0000000000000..7540b29743f60 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/non_max_suppression.h @@ -0,0 +1,34 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" +//#include "core/util/math_cpuonly.h" + +namespace onnxruntime { +namespace contrib { + +template +class NonMaxSuppression final : public OpKernel { + public: + NonMaxSuppression(const OpKernelInfo& info) : OpKernel(info) { + ONNXRUNTIME_ENFORCE(info.GetAttr("max_output_size", &max_output_size_).IsOK()); + ONNXRUNTIME_ENFORCE(info.GetAttr("iou_threshold", &iou_threshold_).IsOK()); + ONNXRUNTIME_ENFORCE(iou_threshold_ >= 0 && iou_threshold_ <= 1, "iou_threshold must be in range [0, 1]"); + ONNXRUNTIME_ENFORCE(info.GetAttr("score_threshold", &score_threshold_).IsOK()); + } + + Status Compute(OpKernelContext* context) const override; + +private: + bool SuppressByIOU(const T* boxes_data, int32_t box_index1, int32_t box_index2) const; + void MaxMin(const T& lhs, const T& rhs, T& min, T& max) const; + +private : int64_t max_output_size_; + float iou_threshold_; + float score_threshold_; +}; +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/non_max_suppression_test.cc b/onnxruntime/test/contrib_ops/non_max_suppression_test.cc new file mode 100644 index 0000000000000..45524fc6626cb --- /dev/null +++ b/onnxruntime/test/contrib_ops/non_max_suppression_test.cc @@ -0,0 +1,205 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime { +namespace test { + +TEST(NonMaxSuppressionOpTest, WithIOUThreshold) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, -0.1f, 1.0f, 0.9f, + 0.0f, 10.0f, 1.0f, 11.0f, + 0.0f, 10.1f, 1.0f, 11.1f, + 0.0f, 100.0f, 1.0f, 101.0f}); + test.AddInput("scores", {6}, {0.9f, 0.75f, 0.6f, 0.95f, 0.5f, 0.3f}); + test.AddAttribute("max_output_size", 3LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {3}, {3L, 0L, 5L}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, WithScoreThreshold) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, -0.1f, 1.0f, 0.9f, + 0.0f, 10.0f, 1.0f, 11.0f, + 0.0f, 10.1f, 1.0f, 11.1f, + 0.0f, 100.0f, 1.0f, 101.0f}); + test.AddInput("scores", {6}, {0.9f, 0.75f, 0.6f, 0.95f, 0.5f, 0.3f}); + test.AddAttribute("max_output_size", 3LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.4f); + test.AddOutput("selected_indices", {2}, {3L, 0L}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, WithScoreThresholdZeroScores) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, -0.1f, 1.0f, 0.9f, + 0.0f, 10.0f, 1.0f, 11.0f, + 0.0f, 10.1f, 1.0f, 11.1f, + 0.0f, 100.0f, 1.0f, 101.0f}); + test.AddInput("scores", {6}, {0.1f, 0.0f, 0.0f, 0.3f, 0.2f, -5.0f}); + test.AddAttribute("max_output_size", 6LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", -3.0f); + test.AddOutput("selected_indices", {2}, {3L, 0L}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, FlippedCoordinates) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {1.0f, 1.0f, 0.0f, 0.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, 0.9f, 1.0f, -0.1f, + 0.0f, 10.0f, 1.0f, 11.0f, + 1.0f, 10.1f, 0.0f, 11.1f, + 1.0f, 101.0f, 0.0f, 100.0f}); + test.AddInput("scores", {6}, {0.9f, 0.75f, 0.6f, 0.95f, 0.5f, 0.3f}); + test.AddAttribute("max_output_size", 3LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {3}, {3L, 0L, 5L}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, SelectTwo) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, -0.1f, 1.0f, 0.9f, + 0.0f, 10.0f, 1.0f, 11.0f, + 0.0f, 10.1f, 1.0f, 11.1f, + 0.0f, 100.0f, 1.0f, 101.0f}); + test.AddInput("scores", {6}, {0.9f, 0.75f, 0.6f, 0.95f, 0.5f, 0.3f}); + test.AddAttribute("max_output_size", 2LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {2}, {3L, 0L}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, SelectThirty) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, -0.1f, 1.0f, 0.9f, + 0.0f, 10.0f, 1.0f, 11.0f, + 0.0f, 10.1f, 1.0f, 11.1f, + 0.0f, 100.0f, 1.0f, 101.0f}); + test.AddInput("scores", {6}, {0.9f, 0.75f, 0.6f, 0.95f, 0.5f, 0.3f}); + test.AddAttribute("max_output_size", 30LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {3}, {3L, 0L, 5L}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, SelectSingleBox) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {1, 4}, + {0.0f, 0.0f, 1.0f, 1.0f}); + test.AddInput("scores", {1}, {0.9f}); + test.AddAttribute("max_output_size", 3LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {1}, {0L}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, SelectFromIdenticalBoxes) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {10, 4}, + {0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.0f, 1.0f, 1.0f, + + 0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.0f, 1.0f, 1.0f}); + test.AddInput("scores", {10}, {0.9f, 0.9f, 0.9f, 0.9f, 0.9f, 0.9f, 0.9f, 0.9f, 0.9f, 0.9f}); + test.AddAttribute("max_output_size", 3LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {1}, {0L}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, InconsistentBoxAndScoreShapes) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, -0.1f, 1.0f, 0.9f, + 0.0f, 10.0f, 1.0f, 11.0f, + 0.0f, 10.1f, 1.0f, 11.1f, + 0.0f, 100.0f, 1.0f, 101.0f}); + test.AddInput("scores", {5}, {0.9f, 0.75f, 0.6f, 0.95f, 0.5f}); + test.AddAttribute("max_output_size", 30LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {1}, {0L}); + test.Run(OpTester::ExpectResult::kExpectFailure, "scores and boxes should have same num_boxes."); +} + +TEST(NonMaxSuppressionOpTest, InvalidIOUThreshold) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {1, 4}, + {0.0f, 0.0f, 1.0f, 1.0f}); + test.AddInput("scores", {1}, {0.9f}); + test.AddAttribute("max_output_size", 3LL); + test.AddAttribute("iou_threshold", 1.2f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {1}, {0L}); + test.Run(OpTester::ExpectResult::kExpectFailure, "iou_threshold must be in range [0, 1]"); +} + +TEST(NonMaxSuppressionOpTest, EmptyInput) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {0, 4}, + {}); + test.AddInput("scores", {0}, {}); + test.AddAttribute("max_output_size", 30LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {0}, {}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, ZeroMaxOutputSize) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {1.0f, 1.0f, 0.0f, 0.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, 0.9f, 1.0f, -0.1f, + 0.0f, 10.0f, 1.0f, 11.0f, + 1.0f, 10.1f, 0.0f, 11.1f, + 1.0f, 101.0f, 0.0f, 100.0f}); + test.AddInput("scores", {6}, {0.9f, 0.75f, 0.6f, 0.95f, 0.5f, 0.3f}); + test.AddAttribute("max_output_size", 0LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddOutput("selected_indices", {0}, {}); + test.Run(); +} + +} // namespace test +} // namespace onnxruntime From d630816142c7d90667467b168ea469b2f079220a Mon Sep 17 00:00:00 2001 From: Hector Li Date: Thu, 29 Nov 2018 11:39:43 -0800 Subject: [PATCH 2/8] fix the warning --- onnxruntime/contrib_ops/cpu/non_max_suppression.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/cpu/non_max_suppression.cc b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc index b77d8880a42a5..2527d4c61fc1d 100644 --- a/onnxruntime/contrib_ops/cpu/non_max_suppression.cc +++ b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc @@ -81,7 +81,7 @@ Status NonMaxSuppression::Compute(OpKernelContext* ctx) const { if (max_output_size_ <= 0 || boxes_dims[0] == 0) { std::vector output_dims(1, 0); TensorShape output_shape(output_dims); - auto output_tensor = ctx->Output(0, output_shape); + ctx->Output(0, output_shape); return Status::OK(); } From 54cc6699db9f303cbe38d6e5477ad605cde3f482 Mon Sep 17 00:00:00 2001 From: Hector Li Date: Thu, 29 Nov 2018 15:01:01 -0800 Subject: [PATCH 3/8] Update schema, add more test case --- onnxruntime/contrib_ops/contrib_ops.cc | 18 +++++++-- .../contrib_ops/cpu/non_max_suppression.cc | 15 ++++++-- .../contrib_ops/cpu/non_max_suppression.h | 7 +++- .../contrib_ops/non_max_suppression_test.cc | 38 +++++++++++++++++++ 4 files changed, 69 insertions(+), 9 deletions(-) diff --git a/onnxruntime/contrib_ops/contrib_ops.cc b/onnxruntime/contrib_ops/contrib_ops.cc index 8a295df3d8b31..b380cb0c5140c 100644 --- a/onnxruntime/contrib_ops/contrib_ops.cc +++ b/onnxruntime/contrib_ops/contrib_ops.cc @@ -374,6 +374,12 @@ The bounding box coordinates corresponding to the selected indices can then be o .Input(0, "boxes", "An input tensor. 2D tensor with shape [num_boxes, 4]", "T1") .Input(1, "scores", "An input tensor. 1D tensor with shape [num_boxes]", "T1") .Output(0, "selected_indices", "selected indices from the boxes tensor.", "T2") + .Output( + 1, + "valid_outputs", + "Optional. A 0-D integer tensor representing the number of valid elements in selected_indices, with the valid elements appearing first.", + "T2", + OpSchema::Optional) .TypeConstraint("T1", {"tensor(float)"}, "Constrain input type to float tensor.") .TypeConstraint("T2", {"tensor(int32)"}, @@ -384,12 +390,18 @@ The bounding box coordinates corresponding to the selected indices can then be o AttributeProto::INT) .Attr( "iou_threshold", - "Float representing the threshold for deciding whether boxes overlap too much with respect to IOU.", - AttributeProto::FLOAT) + "Float representing the threshold for deciding whether boxes overlap too much with respect to IOU. Value range [0, 1]. The default is 0.0", + AttributeProto::FLOAT, + static_cast(0.0f)) .Attr( "score_threshold", "Float tensor representing the threshold for deciding when to remove boxes based on score.", - AttributeProto::FLOAT); + AttributeProto::FLOAT) + .Attr( + "pad_to_max_output_size", + "Optional. 1(true) - the output selected_indices is padded to be of length max_output_size. Defaults to 0(false).", + AttributeProto::INT, + OPTIONAL); } class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SampleOp); diff --git a/onnxruntime/contrib_ops/cpu/non_max_suppression.cc b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc index 2527d4c61fc1d..e720f2605db9d 100644 --- a/onnxruntime/contrib_ops/cpu/non_max_suppression.cc +++ b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc @@ -42,7 +42,7 @@ bool NonMaxSuppression::SuppressByIOU(const T* boxes_data, int32_t box_index1 const T intersection_y_max = std::min(y1_max, y2_max); const T intersection_area = std::max(intersection_x_max - intersection_x_min, static_cast(0.0)) * - std::max(intersection_y_max - intersection_y_min, static_cast(0.0)); + std::max(intersection_y_max - intersection_y_min, static_cast(0.0)); if (intersection_area <= static_cast(0.0)) { return false; @@ -96,7 +96,7 @@ Status NonMaxSuppression::Compute(OpKernelContext* ctx) const { auto LessCompare = [](const ScoreIndexPair& lhs, const ScoreIndexPair& rhs) { return lhs.score < rhs.score; }; - + // Filter by score_threshold_ std::priority_queue, decltype(LessCompare)> sorted_scores_with_index(LessCompare); for (int32_t i = 0; i < num_boxes; ++i) { @@ -129,11 +129,18 @@ Status NonMaxSuppression::Compute(OpKernelContext* ctx) const { } } - std::vector output_dim(1, num_of_selected); + int64_t num_to_copy = pad_to_max_output_size_ == 1 ? max_output_size_ : num_of_selected; + std::vector output_dim(1, num_to_copy); TensorShape output_shape(output_dim); Tensor* selected_indices = ctx->Output(0, output_shape); auto output_data = selected_indices->MutableData(); - memcpy(output_data, selected_index.data(), num_of_selected * sizeof(int32_t)); + memcpy(output_data, selected_index.data(), num_to_copy * sizeof(int32_t)); + + TensorShape valid_outputs_shape(std::vector{1}); + Tensor* valid_outputs = ctx->Output(1, valid_outputs_shape); + if (valid_outputs) { + valid_outputs->MutableData()[0] = num_of_selected; + } return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cpu/non_max_suppression.h b/onnxruntime/contrib_ops/cpu/non_max_suppression.h index 7540b29743f60..74c03dbd7bed2 100644 --- a/onnxruntime/contrib_ops/cpu/non_max_suppression.h +++ b/onnxruntime/contrib_ops/cpu/non_max_suppression.h @@ -13,7 +13,8 @@ namespace contrib { template class NonMaxSuppression final : public OpKernel { public: - NonMaxSuppression(const OpKernelInfo& info) : OpKernel(info) { + NonMaxSuppression(const OpKernelInfo& info) : OpKernel(info), + pad_to_max_output_size_(info.GetAttrOrDefault("pad_to_max_output_size", 0LL)) { ONNXRUNTIME_ENFORCE(info.GetAttr("max_output_size", &max_output_size_).IsOK()); ONNXRUNTIME_ENFORCE(info.GetAttr("iou_threshold", &iou_threshold_).IsOK()); ONNXRUNTIME_ENFORCE(iou_threshold_ >= 0 && iou_threshold_ <= 1, "iou_threshold must be in range [0, 1]"); @@ -26,9 +27,11 @@ class NonMaxSuppression final : public OpKernel { bool SuppressByIOU(const T* boxes_data, int32_t box_index1, int32_t box_index2) const; void MaxMin(const T& lhs, const T& rhs, T& min, T& max) const; -private : int64_t max_output_size_; +private : + int64_t max_output_size_; float iou_threshold_; float score_threshold_; + int64_t pad_to_max_output_size_; }; } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/non_max_suppression_test.cc b/onnxruntime/test/contrib_ops/non_max_suppression_test.cc index 45524fc6626cb..0691154eaa080 100644 --- a/onnxruntime/test/contrib_ops/non_max_suppression_test.cc +++ b/onnxruntime/test/contrib_ops/non_max_suppression_test.cc @@ -201,5 +201,43 @@ TEST(NonMaxSuppressionOpTest, ZeroMaxOutputSize) { test.Run(); } +TEST(NonMaxSuppressionOpTest, PadToFiveOutput) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, -0.1f, 1.0f, 0.9f, + 0.0f, 10.0f, 1.0f, 11.0f, + 0.0f, 10.1f, 1.0f, 11.1f, + 0.0f, 100.0f, 1.0f, 101.0f}); + test.AddInput("scores", {6}, {0.9f, 0.75f, 0.6f, 0.95f, 0.5f, 0.3f}); + test.AddAttribute("max_output_size", 5LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.0f); + test.AddAttribute("pad_to_max_output_size", 1LL); + test.AddOutput("selected_indices", {5}, {3L, 0L, 5L, 0L, 0L}); + test.AddOutput("valid_outputs", {1}, {3L}); + test.Run(); +} + +TEST(NonMaxSuppressionOpTest, WithScoreThresholdPadToSixOutput) { + OpTester test("NonMaxSuppression", 1, onnxruntime::kMSDomain); + test.AddInput("boxes", {6, 4}, + {0.0f, 0.0f, 1.0f, 1.0f, + 0.0f, 0.1f, 1.0f, 1.1f, + 0.0f, -0.1f, 1.0f, 0.9f, + 0.0f, 10.0f, 1.0f, 11.0f, + 0.0f, 10.1f, 1.0f, 11.1f, + 0.0f, 100.0f, 1.0f, 101.0f}); + test.AddInput("scores", {6}, {0.9f, 0.75f, 0.6f, 0.95f, 0.5f, 0.3f}); + test.AddAttribute("max_output_size", 6LL); + test.AddAttribute("iou_threshold", 0.5f); + test.AddAttribute("score_threshold", 0.4f); + test.AddAttribute("pad_to_max_output_size", 1LL); + test.AddOutput("selected_indices", {6}, {3L, 0L, 0L, 0L, 0L, 0L}); + test.AddOutput("valid_outputs", {1}, {2L}); + test.Run(); +} + } // namespace test } // namespace onnxruntime From 3665348462e2cc719e09281fff40172a669dd13a Mon Sep 17 00:00:00 2001 From: Hector Li Date: Thu, 29 Nov 2018 15:44:25 -0800 Subject: [PATCH 4/8] fix build error on Linux --- onnxruntime/contrib_ops/cpu/non_max_suppression.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/cpu/non_max_suppression.h b/onnxruntime/contrib_ops/cpu/non_max_suppression.h index 74c03dbd7bed2..9626a5deb5755 100644 --- a/onnxruntime/contrib_ops/cpu/non_max_suppression.h +++ b/onnxruntime/contrib_ops/cpu/non_max_suppression.h @@ -14,7 +14,7 @@ template class NonMaxSuppression final : public OpKernel { public: NonMaxSuppression(const OpKernelInfo& info) : OpKernel(info), - pad_to_max_output_size_(info.GetAttrOrDefault("pad_to_max_output_size", 0LL)) { + pad_to_max_output_size_(info.GetAttrOrDefault("pad_to_max_output_size", 0)) { ONNXRUNTIME_ENFORCE(info.GetAttr("max_output_size", &max_output_size_).IsOK()); ONNXRUNTIME_ENFORCE(info.GetAttr("iou_threshold", &iou_threshold_).IsOK()); ONNXRUNTIME_ENFORCE(iou_threshold_ >= 0 && iou_threshold_ <= 1, "iou_threshold must be in range [0, 1]"); From 504e0d345a341ab2a3b91a38d45cf12a1bf2a001 Mon Sep 17 00:00:00 2001 From: Hector Li Date: Thu, 29 Nov 2018 16:19:23 -0800 Subject: [PATCH 5/8] update according to review comments --- onnxruntime/contrib_ops/cpu/non_max_suppression.cc | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/non_max_suppression.cc b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc index e720f2605db9d..acf0a88167add 100644 --- a/onnxruntime/contrib_ops/cpu/non_max_suppression.cc +++ b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc @@ -79,8 +79,7 @@ Status NonMaxSuppression::Compute(OpKernelContext* ctx) const { ONNXRUNTIME_RETURN_IF_NOT(scores_shape.GetDims()[0] == num_boxes, "scores and boxes should have same num_boxes."); if (max_output_size_ <= 0 || boxes_dims[0] == 0) { - std::vector output_dims(1, 0); - TensorShape output_shape(output_dims); + TensorShape output_shape({0}); ctx->Output(0, output_shape); return Status::OK(); } @@ -130,13 +129,12 @@ Status NonMaxSuppression::Compute(OpKernelContext* ctx) const { } int64_t num_to_copy = pad_to_max_output_size_ == 1 ? max_output_size_ : num_of_selected; - std::vector output_dim(1, num_to_copy); - TensorShape output_shape(output_dim); + TensorShape output_shape({num_to_copy}); Tensor* selected_indices = ctx->Output(0, output_shape); auto output_data = selected_indices->MutableData(); memcpy(output_data, selected_index.data(), num_to_copy * sizeof(int32_t)); - TensorShape valid_outputs_shape(std::vector{1}); + TensorShape valid_outputs_shape({1}); Tensor* valid_outputs = ctx->Output(1, valid_outputs_shape); if (valid_outputs) { valid_outputs->MutableData()[0] = num_of_selected; From 59cdfd4fa445d7eb44d62dd85b9468134c9af09b Mon Sep 17 00:00:00 2001 From: Hector Li Date: Thu, 29 Nov 2018 16:40:33 -0800 Subject: [PATCH 6/8] fix typo --- onnxruntime/contrib_ops/cpu/non_max_suppression.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/cpu/non_max_suppression.cc b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc index acf0a88167add..b406f79f4ce1c 100644 --- a/onnxruntime/contrib_ops/cpu/non_max_suppression.cc +++ b/onnxruntime/contrib_ops/cpu/non_max_suppression.cc @@ -114,7 +114,7 @@ Status NonMaxSuppression::Compute(OpKernelContext* ctx) const { sorted_scores_with_index.pop(); bool selected = true; - // Check with existing boxes, suppress if exceed the IOU (Intersection Over Union) threadhold + // Check with existing boxes, suppress if exceed the IOU (Intersection Over Union) threshold for (int i = num_of_selected - 1; i >= 0; --i) { if (SuppressByIOU(boxes_data, selected_index[i], next_top_score.index)) { selected = false; From 8fedb2b3a8162f4fadcfe951bbb22b2471fdb8b2 Mon Sep 17 00:00:00 2001 From: Hector Li Date: Fri, 30 Nov 2018 15:58:50 -0800 Subject: [PATCH 7/8] add shape inference --- onnxruntime/contrib_ops/contrib_ops.cc | 33 ++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/contrib_ops.cc b/onnxruntime/contrib_ops/contrib_ops.cc index b380cb0c5140c..fe57df2286e0a 100644 --- a/onnxruntime/contrib_ops/contrib_ops.cc +++ b/onnxruntime/contrib_ops/contrib_ops.cc @@ -12,8 +12,8 @@ namespace onnxruntime { namespace contrib { using ::ONNX_NAMESPACE::AttributeProto; -using ::ONNX_NAMESPACE::OpSchema; using ::ONNX_NAMESPACE::OPTIONAL; +using ::ONNX_NAMESPACE::OpSchema; void RegisterContribSchemas() { ONNX_CONTRIB_OPERATOR_SCHEMA(SampleOp) @@ -376,7 +376,7 @@ The bounding box coordinates corresponding to the selected indices can then be o .Output(0, "selected_indices", "selected indices from the boxes tensor.", "T2") .Output( 1, - "valid_outputs", + "valid_outputs", "Optional. A 0-D integer tensor representing the number of valid elements in selected_indices, with the valid elements appearing first.", "T2", OpSchema::Optional) @@ -400,8 +400,33 @@ The bounding box coordinates corresponding to the selected indices can then be o .Attr( "pad_to_max_output_size", "Optional. 1(true) - the output selected_indices is padded to be of length max_output_size. Defaults to 0(false).", - AttributeProto::INT, - OPTIONAL); + AttributeProto::INT, + OPTIONAL) + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + auto selected_indices_type = ctx.getOutputType(0)->mutable_tensor_type(); + selected_indices_type->set_elem_type(::onnx::TensorProto_DataType::TensorProto_DataType_INT32); + + // If pad_to_max_output_size is set to 1, the output(0) selected_indices will has a fixed shape [max_output_size]. + auto pad_to_max_output_size = ctx.getAttribute("pad_to_max_output_size"); + if (pad_to_max_output_size && 1 == pad_to_max_output_size->i()) { + auto max_output_size = ctx.getAttribute("max_output_size")->i(); + selected_indices_type + ->mutable_shape() + ->add_dim() + ->set_dim_value(max_output_size); + } + + // valid_outputs is optional, shape is [1] + auto num_outputs = ctx.getNumOutputs(); + if (num_outputs > 1) { + auto valid_outputs_shape = ctx.getOutputType(1)->mutable_tensor_type(); + valid_outputs_shape->set_elem_type(::onnx::TensorProto_DataType::TensorProto_DataType_INT32); + valid_outputs_shape + ->mutable_shape() + ->add_dim() + ->set_dim_value(1); + } + }); } class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SampleOp); From d9fabc4abeafb014c69793e1d8c13e0cc7bc7eaa Mon Sep 17 00:00:00 2001 From: Hector Li Date: Fri, 30 Nov 2018 16:01:39 -0800 Subject: [PATCH 8/8] revert a minor change --- onnxruntime/contrib_ops/contrib_ops.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/contrib_ops.cc b/onnxruntime/contrib_ops/contrib_ops.cc index fe57df2286e0a..347da54e8825f 100644 --- a/onnxruntime/contrib_ops/contrib_ops.cc +++ b/onnxruntime/contrib_ops/contrib_ops.cc @@ -12,8 +12,8 @@ namespace onnxruntime { namespace contrib { using ::ONNX_NAMESPACE::AttributeProto; -using ::ONNX_NAMESPACE::OPTIONAL; using ::ONNX_NAMESPACE::OpSchema; +using ::ONNX_NAMESPACE::OPTIONAL; void RegisterContribSchemas() { ONNX_CONTRIB_OPERATOR_SCHEMA(SampleOp)