python实现NMS

python实现NMS

什么叫NMS

NMS即non maximum suppression即非极大抑制。专业的自己看链接,我就只用一句话说我的理解。
就是在多个目标框里,每个目标选择唯一最好的框(红色)。
注意:最好的这个红框可能不是最完美的,但是是可选择的里面最好的。

如何选择最好的,这里就会产生一个新的概念,叫做IoU:交并比(Intersection-over-Union,IoU)

参考博文

目标检测中NMS(非极大抑制)的概念理解

丽丽

以下是我自己写的代码,实现nms功能

def func_IoU(a, b):
    '''
    lx =a[0], ly = a[1] ← 左上角
                            右下角 →  rx =a[2],ry = a[3]
    '''
    # 如果范围出错,退出计算,返回0
    if a[0] > a[2] or a[1] > a[3] or b[0] > b[2] or b[1] > b[3]:
        return 0.0

    # 计算相交面积
    w = min(a[2], b[2]) - max(a[0], b[0])
    h = min(a[3], b[3]) - max(a[1], b[1])
    if w<0 or h <0:
        return 0.0
    in_areas = w * h

    # 计算并集面积
    a_areas = (a[2] - a[0]) * (a[3] - a[1])
    b_areas = (b[2] - b[0]) * (b[3] - b[1])
    uni_areas =a_areas+b_areas-in_areas
    return in_areas/(uni_areas+1e-6)


# 1.选出得分最高的
# 2.和最高得分的面积进行iou计算,大于threshold的去除

def nms(bboxes, confidence, threshold):
    class_list=[]
    conf_list=[]
    if len(bboxes)>0:
        keep_box_list = []
        keep_conf_list= []
        max_num = confidence.index(max(confidence))
        max_box = bboxes.pop(max_num)
        max_conf = confidence.pop(max_num)
        for i in range(len(bboxes)):
            if func_IoU(max_box,bboxes[i])<=threshold:
                keep_box_list.append(bboxes[i])
                keep_conf_list.append(confidence[i])
                print(f'max_box:{max_box},max_conf:{max_conf},\ '
                      f'conf:{confidence[i]}的IoU:{func_IoU(max_box,bboxes[i])}')
        #递归几次,说明有几个prediction窗口
        class_list,conf_list =nms(keep_box_list,keep_conf_list,threshold)
        # 将返回的窗口加入到列表
        class_list.append(max_box)
        conf_list.append(max_conf)
        return class_list,conf_list
    return class_list,conf_list

if __name__ == '__main__':
    bounding=[(187,82,337,317),(150,67,305,282),(246,121,368,304)]
    confidence_score = [0.9,0.75,0.8]
    threshold = 0.6
    picked_boxes,picked_score= nms(bounding,confidence_score,threshold)
    print('最终bbox列表:',picked_boxes)
    print('最终conf分数列表:',picked_score)
    print('threshold:',threshold)
发布了5 篇原创文章 · 获赞 2 · 访问量 102
展开阅读全文

没有更多推荐了,返回首页

©️2019 CSDN 皮肤主题: 1024 设计师: 上身试试

分享到微信朋友圈

×

扫一扫,手机浏览