void DoNMSPerClass(int batch_idx, int class_idx, const float* boxes_data,
                   const float* scores_data, int num_boxes, int q,
                   int num_classes, const int size_per_class,
                   const float score_threshold, const float iou_threshold,
                   std::vector<ResultCandidate>& result_candidate_vec) {
  std::vector<float> class_scores_data;
  class_scores_data.reserve(num_boxes);
  std::vector<float> class_boxes_data;
  class_boxes_data.reserve(num_boxes * 4);

  for (int box_idx = 0; box_idx < num_boxes; ++box_idx) {
    class_scores_data.push_back(scores_data[box_idx * num_classes + class_idx]);
    for (int cid = 0; cid < 4; ++cid) {
      if (q > 1) {
        class_boxes_data.push_back(
            boxes_data[(box_idx * q + class_idx) * 4 + cid]);
      } else {
        class_boxes_data.push_back(boxes_data[box_idx * 4 + cid]);
      }
    }
  }

  // Do NMS, get the candidate indices of form vector<int>
  // Data structure for selection candidate in NMS.
  struct Candidate {
    int box_index;
    float score;
  };
  auto cmp = [](const Candidate bs_i, const Candidate bs_j) {
    return bs_i.score < bs_j.score;
  };
  std::priority_queue<Candidate, std::vector<Candidate>, decltype(cmp)>
      candidate_priority_queue(cmp);
  for (int i = 0; i < num_boxes; ++i) {
    if (class_scores_data[i] > score_threshold) {
      candidate_priority_queue.emplace(Candidate({i, class_scores_data[i]}));
    }
  }

  std::vector<int> selected;
  std::vector<float> selected_boxes;
  Candidate next_candidate;

  // Move class_boxes_data to a tensor
  Eigen::array<Eigen::DenseIndex, 2> boxesShape = {num_boxes, 4};
  typename TTypes<float, 2>::ConstTensor boxes_data_t(class_boxes_data.data(),
                                                      boxesShape);
  float iou;
  while (selected.size() < size_per_class &&
         !candidate_priority_queue.empty()) {
    next_candidate = candidate_priority_queue.top();
    candidate_priority_queue.pop();
    // Overlapping boxes are likely to have similar scores,
    // therefore we iterate through the previously selected boxes backwards
    // in order to see if `next_candidate` should be suppressed.
    bool should_select = true;
    for (int j = selected.size() - 1; j >= 0; --j) {
      iou = IOU<float>(boxes_data_t, next_candidate.box_index, selected[j]);
      if (iou > iou_threshold) {
        should_select = false;
        break;
      }
    }

    if (should_select) {
      // Add the selected box to the result candidate. Sorted by score
      int id = next_candidate.box_index;
      result_candidate_vec[selected.size() + size_per_class * class_idx] = {
          next_candidate.box_index,
          next_candidate.score,
          class_idx,
          {boxes_data_t(id, 0), boxes_data_t(id, 1), boxes_data_t(id, 2),
           boxes_data_t(id, 3)}};
      selected.push_back(next_candidate.box_index);
    }
  }
}
void DoNMSPerClass(int batch_idx, int class_idx, const float* boxes_data,
                   const float* scores_data, int num_boxes, int q,
                   int num_classes, const int size_per_class,
                   const float score_threshold, const float iou_threshold,
                   std::vector<ResultCandidate>& result_candidate_vec) {
  std::vector<float> class_scores_data;
  class_scores_data.reserve(num_boxes);
  std::vector<float> class_boxes_data;
  class_boxes_data.reserve(num_boxes * 4);

  for (int box_idx = 0; box_idx < num_boxes; ++box_idx) {
    class_scores_data.push_back(scores_data[box_idx * num_classes + class_idx]);
    for (int cid = 0; cid < 4; ++cid) {
      if (q > 1) {
        class_boxes_data.push_back(
            boxes_data[(box_idx * q + class_idx) * 4 + cid]);
      } else {
        class_boxes_data.push_back(boxes_data[box_idx * 4 + cid]);
      }
    }
  }

   
   
  struct Candidate {
    int box_index;
    float score;
  };
  auto cmp = [](const Candidate bs_i, const Candidate bs_j) {
    return bs_i.score < bs_j.score;
  };
  std::priority_queue<Candidate, std::vector<Candidate>, decltype(cmp)>
      candidate_priority_queue(cmp);
  for (int i = 0; i < num_boxes; ++i) {
    if (class_scores_data[i] > score_threshold) {
      candidate_priority_queue.emplace(Candidate({i, class_scores_data[i]}));
    }
  }

  std::vector<int> selected;
  std::vector<float> selected_boxes;
  Candidate next_candidate;

   
  Eigen::array<Eigen::DenseIndex, 2> boxesShape = {num_boxes, 4};
  typename TTypes<float, 2>::ConstTensor boxes_data_t(class_boxes_data.data(),
                                                      boxesShape);
  float iou;
  while (selected.size() < size_per_class &&
         !candidate_priority_queue.empty()) {
    next_candidate = candidate_priority_queue.top();
    candidate_priority_queue.pop();
     
     
     
    bool should_select = true;
    for (int j = selected.size() - 1; j >= 0; --j) {
      iou = IOU<float>(boxes_data_t, next_candidate.box_index, selected[j]);
      if (iou > iou_threshold) {
        should_select = false;
        break;
      }
    }

    if (should_select) {
       
      int id = next_candidate.box_index;
      result_candidate_vec[selected.size() + size_per_class * class_idx] = {
          next_candidate.box_index,
          next_candidate.score,
          class_idx,
          {boxes_data_t(id, 0), boxes_data_t(id, 1), boxes_data_t(id, 2),
           boxes_data_t(id, 3)}};
      selected.push_back(next_candidate.box_index);
    }
  }
}