python NMS

NMS

# -*- coding: utf-8 -*-

import cv2
import numpy as np
import random

def py_cpu_nms(dets, thresh):  
    """Pure Python NMS baseline."""  
    x1 = dets[:, 0]
    y1 = dets[:, 1]  
    x2 = dets[:, 2]  
    y2 = dets[:, 3]  
    scores = dets[:, 4] # 获取坐标和分数

    areas = (x2 - x1 + 1) * (y2 - y1 + 1) # 计算每个框的面积 
    order = scores.argsort()[::-1]  # 对分数从高到低排序
    keep = []  
    while order.size > 0:  
        i = order[0]  # 每次都取最高的分数
        keep.append(i)  # 然后保留下来

        xx1 = np.maximum(x1[i], x1[order[1:]])  # 从1开始，第0个就是本身
        yy1 = np.maximum(y1[i], y1[order[1:]])  
        xx2 = np.minimum(x2[i], x2[order[1:]])  
        yy2 = np.minimum(y2[i], y2[order[1:]])  # 计算这个矩形与剩余其他的矩形相交点

        w = np.maximum(0.0, xx2 - xx1 + 1)  
        h = np.maximum(0.0, yy2 - yy1 + 1)  
        inter = w * h  # 计算相交面积

        ovr = inter / (areas[i] + areas[order[1:]] - inter)  # 重叠率，也就是IOU

        inds = np.where(ovr <= thresh)[0] # 重叠率小于阈值的矩形框取出来
        order = order[inds + 1]  # np.where返回的是从零开始的，order是从一开始的，所以要加1
        """
        [0,   1,   2,   3,    4,  ]
        [0.2, 0.4, 0.06, 0.1, 0.24]
        概率由高到低排序
        order [1,   4,    0,   3,   2   ]
        prob  [0.4, 0.24, 0.2, 0.1, 0.06]
        最高概率框与其余框的各种计算（IOU）
        order [1,   4,    0,   3,   2   ]
        prob  [0.4, 0.24, 0.2, 0.1, 0.06] 
        IOU   [1,   0.7,  0.2, 0.6, 0.1]      
        取出IOU<thread(0.5)，得到留下来的框
        order' = order[1:]
        id    [0,   1,   2,   3.   4   ]
        order [1,   4,   0,   3,   2   ]
        id'        [0,   1,   2,   3   ]
        order'     [4,   0,   3,   2   ]
        IOU        [0.7, 0.2, 0.6, 0.1 ]   
        inds' =    [     1,        3   ] # 对应id'
        ids   =    [     2,        4   ] # 对应id
        # 新一轮的开始
        order [4,    2   ]
        prob  [0.24, 0.06]      
        """

    return keep

def plot_rect(dets,img,name):
    font = cv2.FONT_HERSHEY_SIMPLEX
    for point in dets:
        x1,y1,x2,y2,s=point
        x1,y1,x2,y2 = int(x1),int(y1),int(x2),int(y2)
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), thickness=2)
        cv2.rectangle(img, (x1+2, y1+2), (x2, y1+30),(255, 0, 0), thickness=-1)
        cv2.putText(img, str(s), (x1+2, y1+2), font, 0.6, (0, 0, 0), 1)
    cv2.imshow(name,img)
    cv2.waitKey(0)



def main():
    img = np.zeros((1000,1000,3), np.uint8)     
    #fill the image with white  
    img.fill(255)  
    points = []
    for _ in range(10):
        x1 = random.randint(0,500)
        y1 = random.randint(0,500)
        x2 = random.randint(500,1000)
        y2 = random.randint(500,1000)
        s = random.random()
        points.append([x1,y1,x2,y2,s])
    points = np.array(points)
    plot_rect(points,img,'1')
    keep = py_cpu_nms(points,0.5)
    dets = points[keep]
    plot_rect(dets,img,'2')
    cv2.destroyAllWindows()

main()

---

非极大值抑制NMS的python实现