NMS代码

[MiaoRain]

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

import numpy as np
# 输入： dets为候选框; thresh为iou阈值
# 返回值 keep
def nms(dets, thresh):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, -1]
    order = np.argsort(scores)[::-1] # 排序后的索引,由大到小
    areas = (x2-x1) * (y2-y1)
    # print(areas)

    # 最后保留的结果
    keep = []
    while order.size > 0:
        i = order[0] # 取出置信度最高的窗口
        keep.append(i)
        # 计算窗口i 和其他所有窗口 交叠部分面积
        xx1 = np.maximum(x1[i], x1[order[1:]])
        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)
        h = np.maximum(0.0, yy2 - yy1)
        inter = w * h

        # 计算交并比iou
        ovr = inter / ( areas[i] + areas[order[1:] ] - inter )
        print(ovr)
        # inds 为 所有与窗口i的iou值小于threshold值的索引
        # np.where返回值为tuple，*[0]才为真正想要的数据
        inds = np.where(ovr <= thresh)[0]
        print("IOU合并 ", i, order[np.where(ovr > thresh)[0]+1])
        # 因为窗口i代表自身，inds + 1即可对应于order中那些不满足条件的窗口索引位置
        order = order[inds + 1]

    return dets[keep]

import cv2 as cv
import numpy as np
import nms
boxes = [[200, 200, 400, 400, 0.99],
         [220, 220, 420, 420, 0.9],
         [100, 100, 150, 150, 0.82],
         [200, 240, 400, 440, 0.5],
         [150, 250, 300, 400, 0.88]]
boxes = np.array(boxes)
overlap = 0.6 # 阈值
pick = nms.nms(boxes, overlap)
print(pick)

###
canvas = np.zeros((500, 500, 3), dtype='uint8')
color = (0, 255, 255)
for box in boxes:
    x1, y1, x2, y2, _ = map(int, box)
    cv.rectangle(canvas, (x1, y1), (x2, y2), color, 5)
color = (255, 0, 255)
for box in pick:
    # print(box)
    x1, y1, x2, y2, _ = map(int, box)
    cv.rectangle(canvas, (x1, y1), (x2, y2), color, 2)
cv.imwrite("result.png", canvas)
cv.imshow('nms', canvas)
cv.waitKey()