跳转到主内容
新派编程网 - 新潮编程理念,全新技术课堂!

如何用Python实现NMS算法?

大家好,我是林译川,今天我们来聊聊NMS算法的Python实现。NMS,也就是非极大值抑制,是目标检测中常用的一种算法,用于去除重叠的检测框,提高检测的准确性。

首先,我们要了解NMS的原理。简单来说,NMS会根据检测框的得分进行排序,然后选择得分最高的框作为最终结果,并将与之重叠度高的框抑制掉,再从剩余的框中选取得分最高的,重复这个过程,直到没有框可以抑制。

接下来,我们来看看IOU算法,它是计算两个框重叠度的关键。IOU,即交并比,通过计算两个框的交集面积与并集面积的比值来衡量它们的重叠程度。

下面是Python实现NMS的代码,我已经添加了详细的注释,方便大家理解:

import numpy as np

boxes=np.array([[100,100,210,210,0.72],
        [250,250,420,420,0.8],
        [220,220,320,330,0.92],
        [100,100,210,210,0.72],
        [230,240,325,330,0.81],
        [220,230,315,340,0.9]]) 

def py_cpu_nms(dets, thresh):

    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    areas = (y2-y1+1) * (x2-x1+1)
    scores = dets[:,4]
    keep = []
    index = scores.argsort()[::-1]
    while index.size >0:
        i = index[0]       
        keep.append(i)

        x11 = np.maximum(x1[i], x1[index[1:]])    
        y11 = np.maximum(y1[i], y1[index[1:]])
        x22 = np.minimum(x2[i], x2[index[1:]])
        y22 = np.minimum(y2[i], y2[index[1:]])
        

        w = np.maximum(0, x22-x11+1)    
        h = np.maximum(0, y22-y11+1)   
        
        overlaps = w*h
        ious = overlaps / (areas[i]+areas[index[1:]] - overlaps)

        idx = np.where(ious<=thresh)[0]
        index = index[idx+1]   

    return keep

最后,我们来看看效果。通过matplotlib库,我们可以可视化NMS前后的检测结果,红色框表示经过NMS处理后的结果。

import matplotlib.pyplot as plt
def plot_bbox(dets, c='k'):
    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    
    plt.plot([x1,x2], [y1,y1], c)
    plt.plot([x1,x1], [y1,y2], c)
    plt.plot([x1,x2], [y2,y2], c)
    plt.plot([x2,x2], [y1,y2], c)
    plt.title(
                            

相关文章