IOU 和 NMS

IOU：计算交并比

def iou(box1,box2):
    x1min, y1min, x1max, y1max = box1[0], box1[1], box1[2], box1[3]
    x2min, y2min, x2max, y2max = box2[0], box2[1], box2[2], box2[3]
    #计算两个框的面积
    s1 = (y1max - y1min + 1.) * (x1max - x1min + 1.)
    s2 = (y2max - y2min + 1.) * (x2max - x2min + 1.)

    #计算相交部分的坐标
    xmin = max(x1min,x2min)
    ymin = max(y1min,y2min)
    xmax = min(x1max,x2max)
    ymax = min(y1max,y2max)
    print("相交部分：\n",xmin,ymin,xmax,ymax)
    
    # 0是为了防止出现不相交情况
    inter_h = max(ymax - ymin + 1, 0)
    inter_w = max(xmax - xmin + 1, 0)

    intersection = inter_h * inter_w
    union = s1 + s2 - intersection

    #计算iou
    iou = intersection / union
    return iou

box1 = [100,100,200,200]
box2 = [120,120,220,220]  #相交
box3 = [300,300,400,400]  #不相交
box4 = [150,150,150,150]  #b在a里面
print("相交：",iou(box1,box2))
print("不相交：",iou(box1,box3))
print("b在a里面：",iou(box1,box4))

结果：

NMS：在同一个物体上只保留一个边界框

import numpy as np
def non_max_suppress(predicts_dict, threshold=0.2):
    """
    implement non-maximum supression on predict bounding boxes.
    Args:
        predicts_dict: {"stick": [[x1, y1, x2, y2, scores1], [...]]}.
        threshold: iou threshold
    Return:
        predicts_dict processed by non-maximum suppression
    """
    for object_name, bbox in predicts_dict.items():   #对每一个类别的目标分别进行NMS
        bbox_array = np.array(bbox, dtype=np.float)
        # 直接删除置信度小于该阈值的候选框
        con_threshold = 0.4
        bbox_array = bbox_array[bbox_array[:,4] > con_threshold]
 
        ## 获取当前目标类别下所有矩形框（bounding box,下面简称bbx）的坐标和confidence,并计算所有bbx的面积
        x1, y1, x2, y2, scores = bbox_array[:,0], bbox_array[:,1], bbox_array[:,2], bbox_array[:,3], bbox_array[:,4]
       # 计算各类别的每一个边界框的面积
        areas = (x2-x1+1) * (y2-y1+1)
        #print("areas shape = ", areas.shape)
 
        ## 对当前类别下所有的bbx的confidence进行从高到低排序（order保存索引信息）
        order = scores.argsort()[::-1]
        print("order = ", order)
        keep = [] #用来存放最终保留的bbx的索引信息
 
        ## 依次从按confidence从高到低遍历bbx，移除所有与该矩形框的IOU值大于threshold的矩形框
        while order.size > 0:
            i = order[0]
            keep.append(i) #保留当前最大confidence对应的bbx索引
 
            ## 获取所有与当前bbx的交集对应的左上角和右下角坐标，并计算IOU（注意这里是同时计算一个bbx与其他所有bbx的IOU）
            xx1 = np.maximum(x1[i], x1[order[1:]])  #当order.size=1时，下面的计算结果都为np.array([]),不影响最终结果
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])
            print("xx1",xx1)
            print("yy1",yy1)
            # 计算相交面积
            inter = np.maximum(0.0, xx2-xx1+1) * np.maximum(0.0, yy2-yy1+1)
            iou = inter/(areas[i]+areas[order[1:]]-inter)
            print("iou =", iou)
 
            print(np.where(iou<=threshold)) #输出没有被移除的bbx索引（相对于iou向量的索引）
            indexs = np.where(iou<=threshold)[0] + 1 #获取保留下来的索引(因为没有计算与自身的IOU，所以索引相差１，需要加上)
            print(f"indexs = {indexs}  type = {type(indexs)}")
            order = order[indexs] #更新保留下来的索引
            print("order = ", order)
        bbox = bbox_array[keep]
        predicts_dict[object_name] = bbox#.tolist()
        predicts_dict = predicts_dict
    return predicts_dict

演示：

import numpy as np
image = [[40,30,290,290,0.98],
        [80,10,320,310,0.88],
        [10,100,200,300,0.75],
        [10,320,200,490,0.30],
        [490,100,760,410,0.94],
        [440,10,700,260,0.44],
        [600,300,800,500,0.36],
        [550,160,680,310,0.80]]
image_dic = {"dog":image}
# nms删除多余的框
nms_bbox = non_max_suppress(image_dic)
# 选出最终边界框的坐标
dwon_bbox = nms_bbox["dog"][:,:-1].astype(int)
# 在图中显示
import cv2 as cv
image = cv.imread("dog.png")
image2 = cv.imread("dog.png")
cv.imshow("image",image)
# 金毛
cv.rectangle(image,(40,30),(290,290),(0,0,255),5)        #0.98
cv.rectangle(image,(80,10),(320,310),(255,255,255),5)    #0.88
cv.rectangle(image,(10,100),(200,300),(255,0,255),5)     #0.75
cv.rectangle(image,(10,320),(200,490),(0,255,0),5)       #0.30

# # 哈士奇
cv.rectangle(image,(490,100),(760,410),(0,0,255),5)     #0.94
cv.rectangle(image,(440,10),(700,260),(225,255,0),5)    #0.44
cv.rectangle(image,(600,300),(800,500),(0,255,255),5)   #0.36
cv.rectangle(image,(550,160),(680,310),(255,0,0),5)     #0.80

cv.imshow("NO_NMS",image)
for i in dwon_bbox:
    cv.rectangle(image2,(i[0],i[1]),(i[2],i[3]),(0,255,255),5) 
cv.imshow("NMS",image2)
cv.waitKey(0)

结果：

原图：

没经过nms处理前：

经过nms处理后：