Faster RCNN 推理 从头写 java (六) NMS (非最大值抑制)

• 1. 图片预处理
• 2. RPN网络预测
• 3. RPN to ROIs
• 4. Classifier 网络预测
• 5. Classifier网络输出对 ROIs过滤与修正
• 6. NMS (非最大值抑制)
• 7. 坐标转换为原始图片维度

一: 输入输出

输入:

• candidate_bboxes: shape 为 [M, 4]
• candidate_probs: shape 为 [M, ]

输出:

• bboxes_on_resize: VGG16图片输入维度下的bbox, shape为 [M, 4] N 表示未知
• result_probs: 最终boxes的概率, shape为 [M, ] M 表示未知

二: 流程

• 计算每个box的面积
• 对概率值排序
• 提取最大概率的box
• 计算最大概率的box与剩余所有box的面积交集
• 计算overlap
• 过滤掉所有与最大概率的box 的IOU > 0.7
• 过滤后还有剩余box 就跳转到步骤3

三: code by code

提取 [x1, y1, x2, y2]

INDArray x1 = boxes.get(NDArrayIndex.all(), NDArrayIndex.point(0));
INDArray y1 = boxes.get(NDArrayIndex.all(), NDArrayIndex.point(1));
INDArray x2 = boxes.get(NDArrayIndex.all(), NDArrayIndex.point(2));
INDArray y2 = boxes.get(NDArrayIndex.all(), NDArrayIndex.point(3));

计算每个box的面积

INDArray area = x2.sub(x1).mul(y2.sub(y1));

对bbox的概率从小到大排序，返回去下标索引.

List idxs = new ArrayList<>(argsort_v1(probs.toFloatVector()));

所有最终留下的box下标索引.

List pick = new ArrayList<>();

遍历，获取最后一个bbox, 也就是概率最大的box

while (idxs.size() > 0)
{
    int last = idxs.size() - 1;
    // 获取概率最大的下标.
    int i = idxs.get(last);
    pick.add(i);

获取除了最大概率box之外的所有box的索引

SpecifiedIndex exceptLastIndex = new SpecifiedIndex(Ints.toArray(idxs.subList(0, last)));

提取出最大概率bbox之外的所有bbox的矩形框数据
并计算面积交集框

INDArray xx1 = Transforms.max(x1.get(exceptLastIndex), x1.getDouble(i));
INDArray yy1 = Transforms.max(y1.get(exceptLastIndex), y1.getDouble(i));
INDArray xx2 = Transforms.min(x2.get(exceptLastIndex), x2.getDouble(i));
INDArray yy2 = Transforms.min(y2.get(exceptLastIndex), y2.getDouble(i));

约束宽高最小为0

INDArray ww = Transforms.max(xx2.sub(xx1), 0);
INDArray hh = Transforms.max(yy2.sub(yy1), 0);

计算面积交集

INDArray area_intersection = ww.mul(hh);

计算面积并集

INDArray area_union = area.get(exceptLastIndex).add(area.getFloat(i)).sub(area_intersection);

计算最大概率的bbox与其他所有box的overlap, 也叫IOU, 交集面积 / 并集面积

INDArray overlap = area_intersection.div(area_union.add(1e-6));

删除掉与最大bbox 的 IOU 大于阀值的bbox

for (int k = idxs.size() - 1; k >= 0 ; k--)
{
    if (overlap.getFloat(k) > overlap_thresh)
    {
        idxs.remove(k);
    }
}

提取出最终留下的bbox, 并返回函数.

INDArray valid_boxes = boxes.get(valid_index, NDArrayIndex.all()).castTo(DataType.INT);
INDArray valid_probs = probs.get(valid_index);

return Arrays.asList(valid_boxes, valid_probs);