欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页  >  IT编程

IOU与NMS实现代码

程序员文章站 2022-03-03 10:53:53
import numpy as npimport 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...
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()

IOU与NMS实现代码
IOU与NMS实现代码
IOU与NMS实现代码

本文地址:https://blog.csdn.net/weixin_38132153/article/details/107634342