工作生活

Faster RCNN 推理 从头写 java (六) NMS

2019-07-02  本文已影响0人  tabsong

目录:

一: 输入输出

输入:

输出:

二: 流程

三: 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<Integer> idxs = new ArrayList<>(argsort_v1(probs.toFloatVector()));

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

List<Integer> 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);
上一篇 下一篇

猜你喜欢

热点阅读