java 实现非极大值抑制

数据预处理

对于物体检测的结果bbox[][] 和 scores[]，假设输出结果有M个，bbox.shape = [M, 4]，scores.shape =[M*N], N 为检测物体的类的数目。对于单类的物体检测，可以直接使用；若想应用于多类的物体检测，则需要对检测结果做一步处理，遍历scores数组，将每个结果的置信度最大值保存至一个长度为M的一维数组，同时存下对应的类至另一个数组（如第i个bbox对于苹果的置信度是0.3，对于香蕉的置信度是0.7，那就把0.7和香蕉分别保存到两个List里面）。

List<Float> bbox_max_scores = new ArrayList<>();
List<Integer> bbox_max_scores_classes = new ArrayList<>();//0表示apple，1表示bananafor (int i = 0; i < scores.length; i += 2) {if (score[i] > score[i + 1]) {bbox_max_scores_classes.add(0);bbox_max_scores.add(score[i]);} else {bbox_max_scores_classes.add(1);bbox_max_scores.add(score[i + 1]);}}

然后对bbox_max_scores进行非极大值抑制，得到结果后再查找对应下标，得到这个bbox对应的类是苹果还是香蕉。

非极大值抑制

给定阈值CONF_THRESH, 提取出所有置信度大于阈值的bbox的下标和值，按照置信度升序排列，开始：
1.循环取出得分最高的bbox，保存到pickList，计算剩余所有bbox与它的交并比，移除所有iou大于iou_thresh的bbox以及它本身；
2.在剩余的bbox中重复第一步的过程；
3.直到pickList里元素多于NUM_DETECTIONS，或者所有bbox都计算完毕，退出循环。

public List<Integer> single_class_non_max_suppression(float[][] bboxes, List<Float> confidences) {if (bboxes.length == 0)return null;List<Integer> conf_keep_idx = new ArrayList<>();//保存置信度大于CONF_THRESH的元素的下标List<Float> new_confidences = new ArrayList<>();//保存置信度大于CONF_THRESH的元素的值for (int i = 0; i < confidences.size(); i++) {float confidence = confidences.get(i);if (confidence> CONF_THRESH){conf_keep_idx.add(i);new_confidences.add(confidence);}}if (conf_keep_idx.isEmpty())return null;int alive_idx_size = conf_keep_idx.size();List<Idxs> idxsList = new ArrayList<>();//将置信度与下标对应for (int i = 0; i < alive_idx_size; i++) {Idxs idxs = new Idxs(conf_keep_idx.get(i), new_confidences.get(i));idxsList.add(idxs);}//按score升序排列Collections.sort(idxsList);float overlap_xmin, overlap_ymin, overlap_xmax, overlap_ymax;float overlap_w, overlap_h, overlap_area, overlap_ratio;//取出得分最高的bbox，计算剩下的bbox与它的交并比iou，去掉大于iou_thresh的bboxList<Integer> pickList = new ArrayList<>();while (idxsList.size() > 0) {sleep(20);int last = idxsList.size() - 1;if(pickList.size() >= NUM_DETECTIONS)//取置信度最高的NUM_DETECTIONS个结果break;int last_index = idxsList.get(last).getIndex();float last_area = (bboxes[last_index][2] -bboxes[last_index][0]) * (bboxes[last_index][3] -bboxes[last_index][1]);//area=(xmax-xmin)*(ymax-ymin)pickList.add(last_index);List<Idxs> idxs_to_remove = new ArrayList<>();//交并比过大需要移除的bboxfor (int i = 0; i < last; i++) {int i_index = idxsList.get(i).getIndex();overlap_xmin = Math.max(bboxes[last_index][0], bboxes[i_index][0]);overlap_ymin = Math.max(bboxes[last_index][1], bboxes[i_index][1]);overlap_xmax = Math.min(bboxes[last_index][2], bboxes[i_index][2]);overlap_ymax = Math.min(bboxes[last_index][3], bboxes[i_index][3]);overlap_w = Math.max(0, overlap_xmax - overlap_xmin);overlap_h = Math.max(0, overlap_ymax - overlap_ymin);overlap_area = overlap_w * overlap_h;float i_area = (bboxes[i_index][2] -bboxes[i_index][0]) * (bboxes[i_index][3] -bboxes[i_index][1]);overlap_ratio = overlap_area / ( last_area + i_area - overlap_area);//IOUif (overlap_ratio > IOU_THRESH)idxs_to_remove.add(idxsList.get(i));}idxs_to_remove.add(idxsList.get(last));idxsList.removeAll(idxs_to_remove);}return pickList;}

pickList里存放的是满足条件的bbox的下标，得到下标之后再去查找对应bbox的xmin, ymin, xmax, ymax和这个框的物体检测的类(class)和置信度(score)。