non_max_suppression.cc | Explore in Territory

/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/lite/kernels/internal/reference/non_max_suppression.h"

#include <initializer_list>

#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/tensor.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"

namespace tflite {
namespace ops {
namespace builtin {
namespace non_max_suppression {

// Boxes in format [y1, x1, y2, x2]. Shape: [num_boxes, 4]
// Type: Float.
constexpr int kInputTensorBoxes = …;
// Shape: [num_boxes]
// Type: Float.
constexpr int kInputTensorScores = …;
// Max number of boxes to output. Actual output can be smaller.
// The output tensors (indices/scores) are of this length.
// Type: Int32.
constexpr int kInputTensorMaxOutputSize = …;
// Type: Float.
constexpr int kInputTensorIouThreshold = …;
// Type: Float.
constexpr int kInputTensorScoreThreshold = …;
// Only applies to NON_MAX_SUPPRESSION_V5.
// Type: Float.
constexpr int kInputTensorSigma = …;

// Indices of selected boxes. Shape: [num_selected_indices]
// Type: Int32.
constexpr int kNMSOutputTensorSelectedIndices = …;
// Type: Int32.
constexpr int kNMSOutputTensorNumSelectedIndices = …;

// Indices of selected boxes. Shape: [num_selected_indices]
// Type: Int32.
constexpr int kSoftNMSOutputTensorSelectedIndices = …;
// Scores of selected boxes. Shape: [num_selected_indices]
// Type: Float.
constexpr int kSoftNMSOutputTensorSelectedScores = …;
// Type: Int32.
constexpr int kSoftNMSOutputTensorNumSelectedIndices = …;

TfLiteStatus SetTensorSizes(TfLiteContext* context, TfLiteTensor* tensor,
                            std::initializer_list<int> values) { … }

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { … }

// If num_selected_indices < max_output_size, the output tensor can contain
// garbage values initially present in memory. This causes segfault in
// downstream ops such as GATHER, since one of the outputs denotes indices and
// int garbage values can be pretty large. This method zeroes-out the remaining
// values.
// NOTE: We ensure memory being reset is valid, by setting pertinent output
// tensors to max_output_size length in Prepare.
void ResetUnusedElementsToZeroes(const int max_output_size,
                                 const int num_selected_indices,
                                 int* selected_indices,
                                 float* selected_scores) { … }

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { … }
}  // namespace non_max_suppression

TfLiteRegistration* Register_NON_MAX_SUPPRESSION_V4() { … }

TfLiteRegistration* Register_NON_MAX_SUPPRESSION_V5() { … }

}  // namespace builtin
}  // namespace ops
}  // namespace tflite
chromium/third_party/tflite/src/tensorflow/lite/kernels/non_max_suppression.cc
