目标检测nms python代码_[目标检测][python][cpp]非极大值抑制(NMS)算法原理以及CPP实现...

问题描述

在目标检测中，有一个很重要的算法，就是非极大值抑制算法，它本身是一个贪心算法。在多个目标检测预测框结果里找到极大的那个，也即是置信度最高的那个。最近有被问到有关NMS的CPP实现，大概查了一下，大部分都是用python写的，用cpp可能更困难一些。

解决思路

算法原理

输入：包含多个输入框和置信度的N*5矩阵，其中N为检测框的数量，阈值Th

输出：检测物体不重复的框M*5的矩阵，M为图片中目标数量

将检测矩阵中的所有框按照score排序

从头选择最大score的检测框，放入到输出矩阵中，然后遍历剩余的框，并计算与最大框的IOU，当IOU大于Th时，那么该框受到抑制，从输入框中删除。小于阈值则保留。

从输入检测矩阵的剩余的框中，找到下一个最大的框，扔进输出矩阵，重复第二步的操作。

直到输入框为空，即为所有的输入都被处理了，此时输出矩阵里的检测框即为图片中M个对象的检测框。

python实现

import numpy as np

def NMS(dets, thresh):

x1 = dets[:,0]

y1 = dets[:,1]

x2 = dets[:,2]

y2 = dets[:,3]

scores = dets[:,4]

areas = (x2 - x1 + 1) * (y2 - y1 + 1)

order=score.argsort()[::-1]

keep = []

while order.size() > 0:

p = order[0]

keep.append(p)

xx1 = np.maximum(x1[p], x1[1:])

yy1 = np.maximum(y1[p], y1[1:])

xx2 = np.minimum(x2[p], y2[1:])

yy2 = np.minimum(y2[p], y2[1:])

inter = (xx2 - xx1 + 1) * (yy2 - yy1 + 1)

inter = np.maximum(inter, 0)

over = inter / (areas[p] + areas[1:] - inter)

inds = np.where(over < thresh)[0]

order = order[inds + 1]

return dets[keep]

CPP实现

#include

#include

using namespace std;

const int N = 1e6 + 1;

const int M = 82;

class Node

{

public:

double x1, x2, y1, y2;

float score;

double area()

{

double w = max(0.0, x2 - x1 + 1);

double h = max(0.0, y2 - y1 + 1);

return w * h;

}

};

int cmp(const Node A, const Node B)

{

return A.score > B.score;

}

double IoU(Node A, Node B)

{

double x1 = max(A.x1, B.x1);

double y1 = max(A.y1, B.y1);

double x2 = max(A.x2, B.x2);

double y2 = max(A.y2, B.y2);

Node inter = {x1, x2, y1, y2, 0.0};

return inter.area() / (A.area() + B.area() - inter.area());

}

vector NMS(vector dets, float therh)

{

vectorans;

sort(dets.begin(), dets.end(), cmp);

while (dets.size() > 0)

{

Node p = dets[0];

ans.push_back(p);

for (int i=1; i< dets.size(); i++)

{

if (IoU(p, dets[i]) >= therh)

{

dets.erase(dets.begin()+i);

}

}

dets.erase(dets.begin()+0);

}

return ans;

}

int main(int argc, char const *argv[])

{

Node A = {1, 2, 3, 4, 0.6};

Node B = {5, 6, 6, 8, 0.3};

Node C = {1.2, 2.1, 3.1, 4.1, 0.66};

vector dets;

dets.push_back(A);

dets.push_back(B);

dets.push_back(C);

vector ans = NMS(dets, 0.5);

for (int i=0; i

{

cout<

}

system("pause");

return 0;

}

cpp的实现略微麻烦一些，主要是因为需要自己手写结构，并且需要生成排序算法，当然，由于使用了std::vector，在删除时需要移动元素，因此还有很大的可改进空间，比如增加一个flag数组。

其他变体

其他变体比如softNMS是将IoU大于阈值的框的score设定较高的衰减率，这样IoU较大的框的分数会不断下降，直至小于阈值。使用这种方法和NMS的区别在于，NMS直接删除，将对应框的score设置为0，而softNMS则设定一个下降值。