NMS算法的c++实现
算法原理:
先假设有6个候选框,根据分类器类别分类概率做排序,从小到大分别属于车辆的概率分别为A、B、C、D、E、F。
从最大概率矩形框(即面积最大的框)F开始,分别判断A~E与F的重叠度IOU是否大于某个设定的阈值;
假设B、D与F的重叠度超过阈值,那么就扔掉B、D(因为超过阈值,说明D与F或者B与F,已经有很大部分是重叠的,那我们保留面积最大的F即可,其余小面积的B,D就是多余的,用F完全可以表示一个物体了,所以保留F丢掉B,D);并标记第一个矩形框F,是我们保留下来的。
从剩下的矩形框A、C、E中,选择概率最大的E,然后判断E与A、C的重叠度,重叠度大于一定的阈值,那么就扔掉;并标记E是我们保留下来的第二个矩形框。
一直重复这个过程,找到所有曾经被保留下来的矩形框。
代码中定义的box是左上角的x、y以及宽度w、高度h
#include<iostream>
#include<vector>
#include<algorithm>
using namespace std;
typedef struct Bbox{
int x;
int y;
int w;
int h;
float score;
}Bbox;
static bool sort_score(Bbox box1,Bbox box2){
return box1.score > box2.score ? true : false;
}
float iou(Bbox box1,Bbox box2){
int x1 = max(box1.x,box2.x);
int y1 = max(box1.y,box2.y);
int x2 = min(box1.x+box1.w,box2.x+box2.w);
int y2 = min(box1.y+box1.h,box2.y+box2.h);
int w = max(0,x2 - x1);
int h = max(0,y2 - y1);
float over_area = w*h;
return over_area/(box1.w * box1.h + box2.w * box2.h - over_area);
}
vector<Bbox> nms(std::vector<Bbox>&vec_boxs,float threshold){
vector<Bbox>results;
std::sort(vec_boxs.begin(),vec_boxs.end(),sort_score);
while(vec_boxs.size() > 0)
{
results.push_back(vec_boxs[0]);
int index = 1;
while(index < vec_boxs.size()){
float iou_value = iou(vec_boxs[0],vec_boxs[index]);
cout << "iou:" << iou_value << endl;
if(iou_value > threshold)
vec_boxs.erase(vec_boxs.begin() + index);
else
index++;
}
vec_boxs.erase(vec_boxs.begin());
}
return results;
}
int main(){
vector<Bbox> input;
Bbox box1 = {1,1,1,1,0.3};
Bbox box2 = {0,0,2,2,0.4};
//Bbox box1 = {1,3,2,2,0.3}; //iou为负
//Bbox box2 = {0,0,2,2,0.4};
//Bbox box1 = {4,4,2,2,0.3};
//Bbox box2 = {0,0,2,2,0.4};
// Bbox box1 = {4,4,1,1,0.3};
// Bbox box2 = {0,0,2,2,0.4};
// Bbox box1 = {3,3,1,1,0.3};
// Bbox box2 = {0,0,2,2,0.4};
input.push_back(box1);
input.push_back(box2);
vector<Bbox> res;
res = nms(input, 0.2);
for(int i = 0;i < res.size();i++){
printf("%d %d %d %d %f",res[i].x,res[i].y,res[i].w,res[i].h,res[i].score);
cout << endl;
}
return 0;
}