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

Posted wildkid1024

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了[目标检测][python][cpp]非极大值抑制(NMS)算法原理以及CPP实现相关的知识,希望对你有一定的参考价值。

问题描述

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

解决思路

算法原理

输入:包含多个输入框和置信度的N*5矩阵,其中N为检测框的数量,阈值Th
输出:检测物体不重复的框M*5的矩阵,M为图片中目标数量

  1. 将检测矩阵中的所有框按照score排序
  2. 从头选择最大score的检测框,放入到输出矩阵中,然后遍历剩余的框,并计算与最大框的IOU,当IOU大于Th时,那么该框受到抑制,从输入框中删除。小于阈值则保留。
  3. 从输入检测矩阵的剩余的框中,找到下一个最大的框,扔进输出矩阵,重复第二步的操作。
  4. 直到输入框为空,即为所有的输入都被处理了,此时输出矩阵里的检测框即为图片中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数组。

其他变体

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


以上是关于[目标检测][python][cpp]非极大值抑制(NMS)算法原理以及CPP实现的主要内容,如果未能解决你的问题,请参考以下文章

非极大值抑制算法(NMS)的python实现

非极大值抑制算法(NMS)的python实现

目标检测的非最大值抑制-NMS

目标检测 — NMS

非极大值抑制--NMS

非极大值抑制(NMS)算法详解