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表示banana
    for (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的bbox
        List<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<>();//交并比过大需要移除的bbox
            for (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);//IOU
                if (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)。