手写iounms（python实现）

import cv2
import numpy as np


def get_iou(box1,box2,wh=False):
    if wh == False:
        xmin1, ymin1, xmax1, ymax1 = box1[:4]
        xmin2, ymin2, xmax2, ymax2 = box2[:4]
    else:
        xmin1, ymin1 = int(box1[0] - box1[2] / 2.0), int(box1[1] - box1[3] / 2.0)
        xmax1, ymax1 = int(box1[0] + box1[2] / 2.0), int(box1[1] + box1[3] / 2.0)
        xmin2, ymin2 = int(box2[0] - box2[2] / 2.0), int(box2[1] - box2[3] / 2.0)
        xmax2, ymax2 = int(box2[0] + box2[2] / 2.0), int(box2[1] + box2[3] / 2.0)

    #获取交集左上角与右下角
    xx1=max(xmin1,xmin2)
    yy1=max(ymin1,ymin2)
    xx2=min(xmax1,xmax2)
    yy2=min(ymax1,ymax2)

    area1=(xmax1-xmin1)*(ymax1-ymin1)
    area2=(xmax2-xmin1)*(ymax2-ymin2)

    inter_area=max(0,xx2-xx1)*max(0,yy2-yy1)
    iou=inter_area/(area1+area2-inter_area+1e-6)

    return iou


def nms(bboxs,threshold):
    bboxs=sorted(bboxs,key=lambda x:x[4],reverse=True)
    print(bboxs)

    nms_list=[]
    flag=[0]*len(bboxs)#抑制为1不抑制为0

    for i in range(len(bboxs)):
        box_i=bboxs[i]
        if flag[i]==1:
            continue
        nms_list.append(box_i)
        #print("第————————————————",i,"————————————————轮")
        for j in range(i+1,len(bboxs)):
            if flag[j]==1:
                continue
            box_j=bboxs[j]
            iou=get_iou(box_i,box_j)
            #print(box_j,"与",box_i,"\tiou: ",iou)
            if iou>threshold:
                #print("抑制",box_j)
                flag[j]=1
    return nms_list

def vision(bboxs,):
    img = np.zeros((1000, 1000, 3), np.uint8)
    for bbox in bboxs:
        x1, y1, x2, y2 ,conf = bbox
        cv2.rectangle(img, (x1, y1), (x2, y2), (25, 215, 155), 2)

    cv2.imshow('image', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__=="__main__":
    box1 = [13, 22, 268, 367, 0.124648176]
    box2 = [18, 27, 294, 400, 0.35818103]
    box3 = [234, 123, 466, 678, 0.13638769]
    box4 = [600, 700, 750, 900, 0.32311122]
    box5 = [200, 128, 500, 700, 0.2893330]
    box6 = [800, 40, 900, 403, 0.23643229]
    box7 = [810, 30, 912, 420, 0.12314322]
    box_lists = [box1, box2, box3, box4, box5, box6, box7]
    nms_list=nms(box_lists,0.65)
    print(nms_list)
    vision(box_lists)
    vision(nms_list)

运行结果：

[[18, 27, 294, 400, 0.35818103], [600, 700, 750, 900, 0.32311122], [200, 128, 500, 700, 0.289333], [800, 40, 900, 403, 0.23643229], [234, 123, 466, 678, 0.13638769], [13, 22, 268, 367, 0.124648176], [810, 30, 912, 420, 0.12314322]]
第———————————————— 0 ————————————————轮
[600, 700, 750, 900, 0.32311122] 与 [18, 27, 294, 400, 0.35818103] iou: 0.0
[200, 128, 500, 700, 0.289333] 与 [18, 27, 294, 400, 0.35818103] iou: 0.0724133633920755
[800, 40, 900, 403, 0.23643229] 与 [18, 27, 294, 400, 0.35818103] iou: 0.0
[234, 123, 466, 678, 0.13638769] 与 [18, 27, 294, 400, 0.35818103] iou: 0.04961667980210164
[13, 22, 268, 367, 0.124648176] 与 [18, 27, 294, 400, 0.35818103] iou: 0.8157546210021713
抑制 [13, 22, 268, 367, 0.124648176]
[810, 30, 912, 420, 0.12314322] 与 [18, 27, 294, 400, 0.35818103] iou: 0.0
第———————————————— 1 ————————————————轮
[200, 128, 500, 700, 0.289333] 与 [600, 700, 750, 900, 0.32311122] iou: -0.0
[800, 40, 900, 403, 0.23643229] 与 [600, 700, 750, 900, 0.32311122] iou: 0.0
[234, 123, 466, 678, 0.13638769] 与 [600, 700, 750, 900, 0.32311122] iou: -0.0
[810, 30, 912, 420, 0.12314322] 与 [600, 700, 750, 900, 0.32311122] iou: 0.0
第———————————————— 2 ————————————————轮
[800, 40, 900, 403, 0.23643229] 与 [200, 128, 500, 700, 0.289333] iou: 0.0
[234, 123, 466, 678, 0.13638769] 与 [200, 128, 500, 700, 0.289333] iou: 0.6658665135904301
抑制 [234, 123, 466, 678, 0.13638769]
[810, 30, 912, 420, 0.12314322] 与 [200, 128, 500, 700, 0.289333] iou: 0.0
第———————————————— 3 ————————————————轮
[810, 30, 912, 420, 0.12314322] 与 [800, 40, 900, 403, 0.23643229] iou: 0.6905516803912375
抑制 [810, 30, 912, 420, 0.12314322]
[[18, 27, 294, 400, 0.35818103], [600, 700, 750, 900, 0.32311122], [200, 128, 500, 700, 0.289333], [800, 40, 900, 403, 0.23643229]]

可视化