tensorflow/lite/kernels/detection_postprocess.cc
// This function implements a regular version of Non Maximal Suppression (NMS)
// for multiple classes where
// 1) we do NMS separately for each class across all anchors and
// 2) keep only the highest anchor scores across all classes
// 3) The worst runtime of the regular NMS is O(K*N^2)
// where N is the number of anchors and K the number of
// classes.
This function implements a regular version of Non Maximal Suppression (NMS)
for multiple classes where
1) we do NMS separately for each class across all anchors and
2) keep only the highest anchor scores across all classes
3) The worst runtime of the regular NMS is O(K*N^2)
where N is the number of anchors and K the number of
classes.
This function implements a regular version of Non Maximal Suppression ( NMS ) for multiple classes where 1 ) we do NMS separately for each class across all anchors and 2 ) keep only the highest anchor scores across all classes 3 ) The worst runtime of the regular NMS is O ( K * N^2 ) where N is the number of anchors and K the number of classes .
TfLiteStatus NonMaxSuppressionMultiClassRegularHelper(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 num_detections_per_class =
      std::min(op_data->detections_per_class, op_data->max_detections);
  const int max_detections = op_data->max_detections;
  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, num_detections_per_class > 0);
  int sorted_indices_size = 0;
  std::vector<BoxInfo> box_info_after_regular_non_max_suppression(
      max_detections + num_detections_per_class);
  std::vector<int> num_selected(num_classes);
  NMSTaskParam nms_task_param{context,
                              node,
                              op_data,
                              scores,
                              num_classes,
                              num_boxes,
                              label_offset,
                              num_classes_with_background,
                              num_detections_per_class,
                              max_detections,
                              num_selected};
  int num_threads =
      CpuBackendContext::GetFromContext(context)->max_num_threads();
  if (num_threads == 1) {
    TF_LITE_ENSURE_OK(
        context, ComputeNMSResult(nms_task_param,  0,
                                  num_classes - 1, sorted_indices_size,
                                  box_info_after_regular_non_max_suppression));
  } else {
    std::atomic<int> next_col(num_threads);
    std::vector<NonMaxSuppressionWorkerTask> tasks;
    tasks.reserve(num_threads);
    for (int i = 0; i < num_threads; ++i) {
      tasks.emplace_back(
          NonMaxSuppressionWorkerTask(nms_task_param, next_col, i));
    }
    cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
                                    CpuBackendContext::GetFromContext(context));
    for (int j = 0; j < tasks.size(); ++j) {
      if (tasks[j].sorted_indices_size == 0) {
        continue;
      }
      memcpy(&box_info_after_regular_non_max_suppression[sorted_indices_size],
             &tasks[j].sorted_box_info[0],
             sizeof(BoxInfo) * tasks[j].sorted_indices_size);
      InplaceMergeBoxInfo(box_info_after_regular_non_max_suppression,
                          sorted_indices_size,
                          sorted_indices_size + tasks[j].sorted_indices_size);
      sorted_indices_size = std::min(
          sorted_indices_size + tasks[j].sorted_indices_size, max_detections);
    }
  }
  for (int output_box_index = 0; output_box_index < max_detections;
       output_box_index++) {
    if (output_box_index < sorted_indices_size) {
      const int anchor_index = floor(
          box_info_after_regular_non_max_suppression[output_box_index].index /
          num_classes_with_background);
      const int class_index =
          box_info_after_regular_non_max_suppression[output_box_index].index -
          anchor_index * num_classes_with_background - label_offset;
      const float selected_score =
          box_info_after_regular_non_max_suppression[output_box_index].score;
      TF_LITE_ENSURE_EQ(context, detection_boxes->type, kTfLiteFloat32);
      TF_LITE_ENSURE_EQ(context, decoded_boxes->type, kTfLiteFloat32);
      ReInterpretTensor<BoxCornerEncoding*>(detection_boxes)[output_box_index] =
          ReInterpretTensor<const BoxCornerEncoding*>(
              decoded_boxes)[anchor_index];
      GetTensorData<float>(detection_classes)[output_box_index] = class_index;
      GetTensorData<float>(detection_scores)[output_box_index] = selected_score;
    } else {
      TF_LITE_ENSURE_EQ(context, detection_boxes->type, kTfLiteFloat32);
      ReInterpretTensor<BoxCornerEncoding*>(
          detection_boxes)[output_box_index] = {0.0f, 0.0f, 0.0f, 0.0f};
      GetTensorData<float>(detection_classes)[output_box_index] = 0.0f;
      GetTensorData<float>(detection_scores)[output_box_index] = 0.0f;
    }
  }
  GetTensorData<float>(num_detections)[0] = sorted_indices_size;
  box_info_after_regular_non_max_suppression.clear();
  return kTfLiteOk;
}
TfLiteStatus NonMaxSuppressionMultiClassRegularHelper ( 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 num_detections_per_class = std :: min ( op_data -> detections_per_class , op_data -> max_detections ) ; const int max_detections = op_data -> max_detections ; 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 , num_detections_per_class > 0 ) ; int sorted_indices_size = 0 ; std :: vector < BoxInfo > box_info_after_regular_non_max_suppression ( max_detections + num_detections_per_class ) ; std :: vector < int > num_selected ( num_classes ) ; NMSTaskParam nms_task_param { context , node , op_data , scores , num_classes , num_boxes , label_offset , num_classes_with_background , num_detections_per_class , max_detections , num_selected } ; int num_threads = CpuBackendContext :: GetFromContext ( context ) -> max_num_threads ( ) ; if ( num_threads == 1 ) { TF_LITE_ENSURE_OK ( context , ComputeNMSResult ( nms_task_param , 0 , num_classes - 1 , sorted_indices_size , box_info_after_regular_non_max_suppression ) ) ; } else { std :: atomic < int > next_col ( num_threads ) ; std :: vector < NonMaxSuppressionWorkerTask > tasks ; tasks . reserve ( num_threads ) ; for ( int i = 0 ; i < num_threads ; ++ i ) { tasks . emplace_back ( NonMaxSuppressionWorkerTask ( nms_task_param , next_col , i ) ) ; } cpu_backend_threadpool :: Execute ( tasks . size ( ) , tasks . data ( ) , CpuBackendContext :: GetFromContext ( context ) ) ; for ( int j = 0 ; j < tasks . size ( ) ; ++ j ) { if ( tasks [ j ] . sorted_indices_size == 0 ) { continue ; } memcpy ( & box_info_after_regular_non_max_suppression [ sorted_indices_size ] , & tasks [ j ] . sorted_box_info [ 0 ] , sizeof ( BoxInfo ) * tasks [ j ] . sorted_indices_size ) ; InplaceMergeBoxInfo ( box_info_after_regular_non_max_suppression , sorted_indices_size , sorted_indices_size + tasks [ j ] . sorted_indices_size ) ; sorted_indices_size = std :: min ( sorted_indices_size + tasks [ j ] . sorted_indices_size , max_detections ) ; } } for ( int output_box_index = 0 ; output_box_index < max_detections ; output_box_index ++ ) { if ( output_box_index < sorted_indices_size ) { const int anchor_index = floor ( box_info_after_regular_non_max_suppression [ output_box_index ] . index / num_classes_with_background ) ; const int class_index = box_info_after_regular_non_max_suppression [ output_box_index ] . index - anchor_index * num_classes_with_background - label_offset ; const float selected_score = box_info_after_regular_non_max_suppression [ output_box_index ] . score ; TF_LITE_ENSURE_EQ ( context , detection_boxes -> type , kTfLiteFloat32 ) ; TF_LITE_ENSURE_EQ ( context , decoded_boxes -> type , kTfLiteFloat32 ) ; ReInterpretTensor < BoxCornerEncoding * > ( detection_boxes ) [ output_box_index ] = ReInterpretTensor < const BoxCornerEncoding * > ( decoded_boxes ) [ anchor_index ] ; GetTensorData < float > ( detection_classes ) [ output_box_index ] = class_index ; GetTensorData < float > ( detection_scores ) [ output_box_index ] = selected_score ; } else { TF_LITE_ENSURE_EQ ( context , detection_boxes -> type , kTfLiteFloat32 ) ; ReInterpretTensor < BoxCornerEncoding * > ( detection_boxes ) [ output_box_index ] = { 0.0f , 0.0f , 0.0f , 0.0f } ; GetTensorData < float > ( detection_classes ) [ output_box_index ] = 0.0f ; GetTensorData < float > ( detection_scores ) [ output_box_index ] = 0.0f ; } } GetTensorData < float > ( num_detections ) [ 0 ] = sorted_indices_size ; box_info_after_regular_non_max_suppression . clear ( ) ; return kTfLiteOk ; }
This function implements a regular version of Non Maximal Suppression (NMS)
for multiple classes where
1) we do NMS separately for each class across all anchors and
2) keep only the highest anchor scores across all classes
3) The worst runtime of the regular NMS is O(K*N^2)
where N is the number of anchors and K the number of
classes.
This function implements a regular version of Non Maximal Suppression ( NMS ) for multiple classes where 1 ) we do NMS separately for each class across all anchors and 2 ) keep only the highest anchor scores across all classes 3 ) The worst runtime of the regular NMS is O ( K * N^2 ) where N is the number of anchors and K the number of classes .
// The row index offset is 1 if background class is included and 0 otherwise. // For each class, perform non-max suppression. /* col_begin= */ // Merge results from tasks. // Allocate output tensors // detection_boxes // detection_classes // detection_scores // detection_classes // detection_scores