tensorflow/lite/kernels/detection_postprocess.cc
// This function implements a fast version of Non Maximal Suppression for
// multiple classes where
// 1) we keep the top-k scores for each anchor and
// 2) during NMS, each anchor only uses the highest class score for sorting.
// 3) Compared to standard NMS, the worst runtime of this version is O(N^2)
// instead of O(KN^2) where N is the number of anchors and K the number of
// classes.
This function implements a fast version of Non Maximal Suppression for
multiple classes where
1) we keep the top-k scores for each anchor and
2) during NMS, each anchor only uses the highest class score for sorting.
3) Compared to standard NMS, the worst runtime of this version is O(N^2)
instead of O(KN^2) where N is the number of anchors and K the number of
classes.
This function implements a fast version of Non Maximal Suppression for multiple classes where 1 ) we keep the top - k scores for each anchor and 2 ) during NMS each anchor only uses the highest class score for sorting . 3 ) Compared to standard NMS the worst runtime of this version is O ( N^2 ) instead of O ( KN^2 ) where N is the number of anchors and K the number of classes .
TfLiteStatus NonMaxSuppressionMultiClassFastHelper(TfLiteContext* context,
                                                   TfLiteNode* node,
                                                   OpData* op_data,
                                                   const float* scores) {
  const TfLiteTensor* input_box_encodings;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputTensorBoxEncodings,
                                 &input_box_encodings));
  const TfLiteTensor* input_class_predictions;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputTensorClassPredictions,
                                 &input_class_predictions));
  const TfLiteTensor* decoded_boxes =
      &context->tensors[op_data->decoded_boxes_index];
  TfLiteTensor* detection_boxes;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensorDetectionBoxes,
                                  &detection_boxes));
  TfLiteTensor* detection_classes;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensorDetectionClasses,
                                  &detection_classes));
  TfLiteTensor* detection_scores;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensorDetectionScores,
                                  &detection_scores));
  TfLiteTensor* num_detections;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensorNumDetections,
                                  &num_detections));
  const int num_boxes = input_box_encodings->dims->data[1];
  const int num_classes = op_data->num_classes;
  const int max_categories_per_anchor = op_data->max_classes_per_detection;
  const int num_classes_with_background =
      input_class_predictions->dims->data[2];
  int label_offset = num_classes_with_background - num_classes;
  TF_LITE_ENSURE(context, (max_categories_per_anchor > 0));
  const int num_categories_per_anchor =
      std::min(max_categories_per_anchor, num_classes);
  std::vector<float> max_scores;
  max_scores.resize(num_boxes);
  std::vector<int> sorted_class_indices;
  sorted_class_indices.resize(num_boxes * num_classes);
  for (int row = 0; row < num_boxes; row++) {
    const float* box_scores =
        scores + row * num_classes_with_background + label_offset;
    int* class_indices = sorted_class_indices.data() + row * num_classes;
    DecreasingPartialArgSort(box_scores, num_classes, num_categories_per_anchor,
                             class_indices);
    max_scores[row] = box_scores[class_indices[0]];
  }
  std::vector<int> selected;
  TF_LITE_ENSURE_STATUS(NonMaxSuppressionSingleClassHelper(
      context, node, op_data, max_scores, op_data->max_detections, &selected));
  int output_box_index = 0;
  for (const auto& selected_index : selected) {
    const float* box_scores =
        scores + selected_index * num_classes_with_background + label_offset;
    const int* class_indices =
        sorted_class_indices.data() + selected_index * num_classes;
    for (int col = 0; col < num_categories_per_anchor; ++col) {
      int box_offset = max_categories_per_anchor * output_box_index + col;
      TF_LITE_ENSURE_EQ(context, detection_boxes->type, kTfLiteFloat32);
      TF_LITE_ENSURE_EQ(context, decoded_boxes->type, kTfLiteFloat32);
      ReInterpretTensor<BoxCornerEncoding*>(detection_boxes)[box_offset] =
          ReInterpretTensor<const BoxCornerEncoding*>(
              decoded_boxes)[selected_index];
      GetTensorData<float>(detection_classes)[box_offset] = class_indices[col];
      GetTensorData<float>(detection_scores)[box_offset] =
          box_scores[class_indices[col]];
    }
    output_box_index++;
  }
  GetTensorData<float>(num_detections)[0] = output_box_index;
  return kTfLiteOk;
}
TfLiteStatus NonMaxSuppressionMultiClassFastHelper ( TfLiteContext * context , TfLiteNode * node , OpData * op_data , const float * scores ) { const TfLiteTensor * input_box_encodings ; TF_LITE_ENSURE_OK ( context , GetInputSafe ( context , node , kInputTensorBoxEncodings , & input_box_encodings ) ) ; const TfLiteTensor * input_class_predictions ; TF_LITE_ENSURE_OK ( context , GetInputSafe ( context , node , kInputTensorClassPredictions , & input_class_predictions ) ) ; const TfLiteTensor * decoded_boxes = & context -> tensors [ op_data -> decoded_boxes_index ] ; TfLiteTensor * detection_boxes ; TF_LITE_ENSURE_OK ( context , GetOutputSafe ( context , node , kOutputTensorDetectionBoxes , & detection_boxes ) ) ; TfLiteTensor * detection_classes ; TF_LITE_ENSURE_OK ( context , GetOutputSafe ( context , node , kOutputTensorDetectionClasses , & detection_classes ) ) ; TfLiteTensor * detection_scores ; TF_LITE_ENSURE_OK ( context , GetOutputSafe ( context , node , kOutputTensorDetectionScores , & detection_scores ) ) ; TfLiteTensor * num_detections ; TF_LITE_ENSURE_OK ( context , GetOutputSafe ( context , node , kOutputTensorNumDetections , & num_detections ) ) ; const int num_boxes = input_box_encodings -> dims -> data [ 1 ] ; const int num_classes = op_data -> num_classes ; const int max_categories_per_anchor = op_data -> max_classes_per_detection ; const int num_classes_with_background = input_class_predictions -> dims -> data [ 2 ] ; int label_offset = num_classes_with_background - num_classes ; TF_LITE_ENSURE ( context , ( max_categories_per_anchor > 0 ) ) ; const int num_categories_per_anchor = std :: min ( max_categories_per_anchor , num_classes ) ; std :: vector < float > max_scores ; max_scores . resize ( num_boxes ) ; std :: vector < int > sorted_class_indices ; sorted_class_indices . resize ( num_boxes * num_classes ) ; for ( int row = 0 ; row < num_boxes ; row ++ ) { const float * box_scores = scores + row * num_classes_with_background + label_offset ; int * class_indices = sorted_class_indices . data ( ) + row * num_classes ; DecreasingPartialArgSort ( box_scores , num_classes , num_categories_per_anchor , class_indices ) ; max_scores [ row ] = box_scores [ class_indices [ 0 ] ] ; } std :: vector < int > selected ; TF_LITE_ENSURE_STATUS ( NonMaxSuppressionSingleClassHelper ( context , node , op_data , max_scores , op_data -> max_detections , & selected ) ) ; int output_box_index = 0 ; for ( const auto & selected_index : selected ) { const float * box_scores = scores + selected_index * num_classes_with_background + label_offset ; const int * class_indices = sorted_class_indices . data ( ) + selected_index * num_classes ; for ( int col = 0 ; col < num_categories_per_anchor ; ++ col ) { int box_offset = max_categories_per_anchor * output_box_index + col ; TF_LITE_ENSURE_EQ ( context , detection_boxes -> type , kTfLiteFloat32 ) ; TF_LITE_ENSURE_EQ ( context , decoded_boxes -> type , kTfLiteFloat32 ) ; ReInterpretTensor < BoxCornerEncoding * > ( detection_boxes ) [ box_offset ] = ReInterpretTensor < const BoxCornerEncoding * > ( decoded_boxes ) [ selected_index ] ; GetTensorData < float > ( detection_classes ) [ box_offset ] = class_indices [ col ] ; GetTensorData < float > ( detection_scores ) [ box_offset ] = box_scores [ class_indices [ col ] ] ; } output_box_index ++ ; } GetTensorData < float > ( num_detections ) [ 0 ] = output_box_index ; return kTfLiteOk ; }
This function implements a fast version of Non Maximal Suppression for
multiple classes where
1) we keep the top-k scores for each anchor and
2) during NMS, each anchor only uses the highest class score for sorting.
This function implements a fast version of Non Maximal Suppression for multiple classes where 1 ) we keep the top - k scores for each anchor and 2 ) during NMS each anchor only uses the highest class score for sorting .
// The row index offset is 1 if background class is included and 0 otherwise. // Perform non-maximal suppression on max scores // Allocate output tensors // detection_boxes // detection_classes // detection_scores