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

2023年4月6日上午12:28 • 目标检测

问题描述

在目标检测中，有一个很重要的算法，就是非极大值抑制算法，它本身是一个贪心算法。在多个目标检测预测框结果里找到极大的那个，也即是置信度最高的那个。最近有被问到有关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 <bits/stdc++.h>
#include <cmath>
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<Node> NMS(vector<Node> dets, float therh)
{
    vector<Node>ans;
    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<Node> dets;
    dets.push_back(A);
    dets.push_back(B);
    dets.push_back(C);
    vector<Node> ans = NMS(dets, 0.5);
    for (int i=0; i<ans.size(); i++)
    {
        cout<<ans[i].area()<<endl;
    }
    system("pause");
    return 0;
}

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