- Notifications
You must be signed in to change notification settings - Fork1
Open
Description
class TreeNode: def __init__(self, x): self.val = x self.left = None self.right = Noneimport numpy as np# 输入: dets为候选框; thresh为iou阈值# 返回值 keepdef 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 cvimport numpy as npimport nmsboxes = [[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()Metadata
Metadata
Assignees
Labels
No labels