NMS实现（C++）

思路

1.对所有Pred框依据置信度排序
2.遍历所有Pred框
3.计算Pred当前框P(i)和其余候选框P(j)的IOU
4.如果IOU>thresh，从预测框集合中剔除P(j)

代码

class Rect
{
public:
    Rect() {};
    //Rect(int x, int y, int w, int h, double score) :x(x), y(y), w(w), h(h), score(score) {};
    Rect(int x, int y, int w, int h, double score)
    {
        this->x = x;
        this->y = y;
        this->w = w;
        this->h = h;
        this->score = score;
    };
    Rect(const Rect& src)
    {
        this->x = src.x;
        this->y = src.y;
        this->w = src.w;
        this->h = src.h;
        this->score = src.score;
    };
    Rect& operator=(const Rect& src)
    {
        this->x = src.x;
        this->y = src.y;
        this->w = src.w;
        this->h = src.h;
        this->score = src.score;
        return *this;
    };
    ~Rect() {};
    int x = -1;
    int y = -1;
    int w = -1;
    int h = -1;
    double score = 0.0;
};

double iou(Rect& r1, Rect& r2)
{
    int i_x1 = std::max(r1.x, r2.x);
    int i_y1 = std::max(r1.y, r2.y);
    int i_x2 = std::min(r1.x + r1.w, r2.x + r2.w);
    int i_y2 = std::min(r1.y + r1.h, r2.y + r2.h);
    if ((i_x2 - i_x1) < 0 || (i_y2 - i_y1) < 0) return 0;
    double i_area = (i_x2 - i_x1) * (i_y2 - i_y1);
    double iou_score = i_area/(r1.w*r1.h + r2.w*r2.h - i_area);
    return iou_score;
};

bool compare(Rect& r1, Rect& r2)
{
    if (r1.score > r2.score) return true;
    return false;
}
void nms(std::vector<Rect>& src, double iou_thresh)
{
    std::sort(src.begin(), src.end() , compare);
    int i = 0, j = 0;
    while( i < src.size())
    {
        j = i+1;
        while (j < src.size())
        {
            float iou_score = iou(src[i], src[j]);
            std::cout << src[i].x << "," << src[i].y << "," << src[i].w << "," << src[i].h << "," << std::to_string(src[i].score) << std::endl;
            std::cout << src[j].x << "," << src[j].y << "," << src[j].w << "," << src[j].h << "," << std::to_string(src[j].score) << std::endl;
            std::cout << iou_score << std::endl;
            if (iou_score > iou_thresh)
            {
                src.erase(src.begin()+j);
                j--;
            }
            j++;
        }
        i++;
    }
}