IOU与NMS实现代码

import numpy as np
import cv2

# A[0-3]分别代表左上角横坐标，左上角纵坐标，右下角横坐标，右下角纵坐标
def IOU(A, B):
	#交集框的左上角，右下角	
    xA = max(A[0], B[0])
    yA = max(A[1], B[1])
    xB = min(A[2], B[2])
    yB = min(A[3], B[3])

    areaA = (A[2] - A[0] + 1) * (A[3] - A[1] + 1)  # 1代表面积要有一个最小值1
    areaB = (B[2] - B[0] + 1) * (B[3] - B[1] + 1)
    #交集，并集
    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)  # 0代表最小是0，不可以出现负数
    iou = interArea/float(areaA + areaB - interArea)

    return iou

# img = np.zeros((400, 400, 3))
# img.fill(255)
# 
# 
# A = [30, 30, 130, 130]
# B = [50, 50, 150, 150]
# # cv2.rectangle(图片,(左上角横坐标,左上角纵坐标),(右下角横坐标,右下角纵坐标),框颜色(B,G,R), 框的宽度):画框
# cv2.rectangle(img, (A[0], A[1]), (A[2], A[3]), (0,255,0), 1)
# cv2.rectangle(img, (B[0], B[1]), (B[2], B[3]), (255,255,0), 1)
# 
# iou = IOU(A, B)
# # cv2.putText(画布，文字，位置，字体，大小，颜色，文字粗细)
# cv2.putText(img, "IOU = %.2f" % iou, (70, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)
# 
# cv2.imshow("img", img)
# cv2.waitKey()


def NMS(dects, thresh):
    sroce = dects[:, 4]
    order = sroce.argsort()[::-1]  # sroce从小到大排序的索引号
    keep = []  # 保存符合条件的index
    while order.size > 0:
        i = order[0]  # 每次选择分数最高的
        keep.append(i)
        iou = []
        for j in range(1,order.size):
            a = IOU(dects[i],dects[order[j]])  # 得分最高的与其他框进行IOU
            iou.append(a)
        iou = np.array(iou)
        index = np.where(iou <= thresh)[0]  # 返回一个元组，[0]是数据，[1]是dtype，大于阈值的舍弃
        order = order[index + 1]  # 新生成的order是小于阈值的位置，并排除第一个位置，例如order=[0,2,1,3],index=[0,1],则新生成的order=[2,1]
    return keep

dects=np.array([[40, 40, 130, 130,0.7],[50, 50, 150, 150,0.9],[60, 60, 160, 160,0.2],[80, 80, 180, 180,0.3]])

img = np.zeros((400, 400, 3))  # 背景大小
img.fill(255)

# dects1 = dects[0,0:4].astype(int)
# dects2 = dects[1,0:4].astype(int)
# dects3 = dects[2,0:4].astype(int)
# dects4 = dects[3,0:4].astype(int)
# # cv2.rectangle(图片,(左上角横坐标,左上角纵坐标),(右下角横坐标,右下角纵坐标),框颜色(B,G,R), 框的宽度):画框
# cv2.rectangle(img, (dects1[0], dects1[1]), (dects1[2], dects1[3]), (0,255,0), 1)
# cv2.rectangle(img, (dects2[0], dects2[1]), (dects2[2], dects2[3]), (255,0,0), 1)
# cv2.rectangle(img, (dects3[0], dects3[1]), (dects3[2], dects3[3]), (255,255,0), 1)
# cv2.rectangle(img, (dects4[0], dects4[1]), (dects4[2], dects4[3]), (255,0,255), 1)

# cv2.imshow("img1", img)
# cv2.waitKey()

order = NMS(dects,0.5)
dects = dects.astype(int)
print(order)
for i in order:
    cv2.rectangle(img, (dects[i,0], dects[i,1]), (dects[i,2], dects[i,3]), (255, 255, 0), 1)

cv2.imshow("img2",img)
cv2.waitKey()