C++(cv::mat)单类别和多类别nms

1.单类别nms

auto nms(const cv::Mat& boxes, float iou_thres){
  // 按照置信度降序排列所有的 box
  int n = boxes.rows;
  cv::Mat sorted_indices;
  // 获取排序后的索引
  cv::sortIdx(boxes.col(4), sorted_indices, cv::SORT_EVERY_COLUMN + cv::SORT_DESCENDING);
  // 根据排序后的索引对原始数组进行重排
  cv::Mat sorted_boxes(boxes.size(), boxes.type());
  for (int i = 0; i < sorted_indices.rows; i++) {
    //std::cout<<sorted_indices.at<int>(i)<<std::endl;
    int idx = sorted_indices.at<int>(i);
    boxes.row(idx).copyTo(sorted_boxes.row(i));
}

  // 定义一个向量，用于存储保留下来的 box 的索引值
  std::vector<int> keep_indices;

  // 定义一个向量，用于存储每个 box 的面积
  std::vector<float> area(n);
  for (int i = 0; i < n; ++i) {
      float x1 = sorted_boxes.at<float>(i, 0);
      float y1 = sorted_boxes.at<float>(i, 1);
      float x2 = sorted_boxes.at<float>(i, 2);
      float y2 = sorted_boxes.at<float>(i, 3);
      area[i] = (x2 - x1 + 1) * (y2 - y1 + 1);

      bool keep = true;
      for (int j = 0; j < keep_indices.size(); ++j) {
          int k = keep_indices[j];
          float xx1 = std::max(x1, sorted_boxes.at<float>(k, 0));
          float yy1 = std::max(y1, sorted_boxes.at<float>(k, 1));
          float xx2 = std::min(x2, sorted_boxes.at<float>(k, 2));
          float yy2 = std::min(y2, sorted_boxes.at<float>(k, 3));
          float w = std::max(0.0f, xx2 - xx1 + 1);
          float h = std::max(0.0f, yy2 - yy1 + 1);
          float overlap = w * h / (area[i] + area[k] - w * h);//两框交并比计算
          if (overlap >= iou_thres) {
              keep = false;
              break;
          }
      }
      if (keep) {
          keep_indices.push_back(i);
      }
  }
  
  // 根据保留下来的 box 的索引值，从原始的 box 信息中提取出保留下来的 box，并返回结果
  //cv::Mat kept_boxes(keep_indices.size(), 6, CV_32F);
  std::vector<std::vector<float>> kept_boxes(keep_indices.size(), std::vector<float>(6, 0));
  for (int i = 0; i < keep_indices.size(); ++i) {
      int k = keep_indices[i];
      //根据你的后处理自行返回结果：std::vector或cv::Mat
      //注释的为cv::Mat
      // kept_boxes.at<float>(i, 0) = sorted_boxes.at<float>(k, 0);
      // kept_boxes.at<float>(i, 1) = sorted_boxes.at<float>(k, 1);
      // kept_boxes.at<float>(i, 2) = sorted_boxes.at<float>(k, 2);
      // kept_boxes.at<float>(i, 3) = sorted_boxes.at<float>(k, 3);
      // kept_boxes.at<float>(i, 4) = sorted_boxes.at<float>(k, 4);
      // kept_boxes.at<float>(i, 5) = sorted_boxes.at<float>(k, 5);//classes 类别维度
      kept_boxes[i][0] = sorted_boxes.at<float>(k, 0);
      kept_boxes[i][1] = sorted_boxes.at<float>(k, 1);
      kept_boxes[i][2] = sorted_boxes.at<float>(k, 2);
      
      kept_boxes[i][3] = sorted_boxes.at<float>(k, 3);
      kept_boxes[i][4] = sorted_boxes.at<float>(k, 4);
      kept_boxes[i][5]= sorted_boxes.at<float>(k, 5);//classes
      
      
  }
  return kept_boxes;
}

2.多类别的nms

先把得出来的结果输出框，按类别进行划分

    //按类别拆分
    std::vector<cv::Mat> class_box_map;
    for(int c = 0; c < num_class; c++){
      cv::Mat classic_filter_res;
      for (int row = 0; row< filter_new_res.rows; row++) {
        if(static_cast<int>(filter_new_res.at<float>(row, 5)) == c){
          cv::Mat row_mat = filter_new_res.rowRange(row, row + 1);
          classic_filter_res.push_back(row_mat);
        }
      }
      class_box_map.push_back(classic_filter_res.clone());
    }

然后依次将每一类送入上文标题1的nms

for(int j = 0; j < class_box_map.size(); j++){
      if(!class_box_map[j].empty()){
        cv:: Mat copy_mat(class_box_map[j].size(), class_box_map[j].type());
        class_box_map[j].copyTo(copy_mat);//可能有的代码没必要加，可自行删除
        //std::cout<<"copy_mat: "<<copy_mat.row(0)<<std::endl;

        std::vector<std::vector<float>> final_boxes = nms(copy_mat, iou_thres);//nms