import numpy as np class Bounding_box: def __init__(self, x1, y1, x2, y2, score): self.x1 = x1 self.y1 = y1 self.x2 = x2 self.y2 = y2 self.score = score def get_iou(boxa, boxb): max_x = max(boxa.x1, boxb.x1) max_y = max(boxa.y1, boxb.y1) min_x = min(boxa.x2, boxb.x2) min_y = min(boxa.y2, boxb.y2) if min_x <= max_x or min_y <= max_y: return 0 area_i = (min_x - max_x) * (min_y - max_y) area_a = (boxa.x2 - boxa.x1) * (boxa.y2 - boxa.y1) area_b = (boxb.x2 - boxb.x1) * (boxb.y2 - boxb.y1) area_u = area_a + area_b - area_i return float(area_i) / float(area_u) def NMS(box_lists, k): box_lists = sorted(box_lists, key=lambda x: x.score, reverse=True) NMS_lists = [box_lists[0]] temp_lists = [] for i in range(k): for j in range(1, len(box_lists)): iou = get_iou(NMS_lists[i], box_lists[j]) if iou < 0.7: temp_lists.append(box_lists[j]) if len(temp_lists) == 0: return NMS_lists box_lists = temp_lists temp_lists = [] NMS_lists.append(box_lists[0]) return NMS_lists box1 = Bounding_box(13, 22, 268, 367, 0.124648176) box2 = Bounding_box(18, 27, 294, 400, 0.35818103) box3 = Bounding_box(234, 123, 466, 678, 0.13638769) box_lists = [box1, box2, box3] NMS_list = NMS(box_lists, 2) print(NMS_list) print(NMS_list[0].x1) def compute_iou_numpy(box1, box2, wh=False): """ compute the iou of two boxes. Args: box1, box2: [xmin, ymin, xmax, ymax] (wh=False) or [xcenter, ycenter, w, h] (wh=True) wh: the format of coordinate. Return: iou: iou of box1 and box2. """ if wh == False: xmin1, ymin1, xmax1, ymax1 = box1 xmin2, ymin2, xmax2, ymax2 = box2 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) ## 获取矩形框交集对应的左上角和右下角的坐标(intersection) xx1 = np.max([xmin1, xmin2]) yy1 = np.max([ymin1, ymin2]) xx2 = np.min([xmax1, xmax2]) yy2 = np.min([ymax1, ymax2]) ## 计算两个矩形框面积 area1 = (xmax1 - xmin1) * (ymax1 - ymin1) area2 = (xmax2 - xmin2) * (ymax2 - ymin2) inter_area = (np.max([0, xx2 - xx1])) * (np.max([0, yy2 - yy1])) iou = inter_area / (area1 + area2 - inter_area + 1e-6) return iou # method1 def calculateIoU(candidateBound, groundTruthBound): ''' :param candidateBound: (xmin,ymin,xmax,ymax) :param groundTruthBound: (xmin,ymin,xmax,ymax) :return: ''' cx1 = candidateBound[0] cy1 = candidateBound[1] cx2 = candidateBound[2] cy2 = candidateBound[3] gx1 = groundTruthBound[0] gy1 = groundTruthBound[1] gx2 = groundTruthBound[2] gy2 = groundTruthBound[3] carea = (cx2 - cx1) * (cy2 - cy1) # C的面积 garea = (gx2 - gx1) * (gy2 - gy1) # G的面积 x1 = max(cx1, gx1) y1 = max(cy1, gy1) x2 = min(cx2, gx2) y2 = min(cy2, gy2) w = max(0, x2 - x1) h = max(0, y2 - y1) area = w * h # C∩G的面积 iou = area / (carea + garea - area) return iou # method2 def compute_iou(rec1, rec2): """ computing IoU :param rec1: (xmin,ymin,xmax,ymax), which reflects (left,top, right,bottom) :param rec2: (xmin,ymin,xmax,ymax) :return: scala value of IoU """ # computing area of each rectangles S_rec1 = (rec1[3] - rec1[1]) * (rec1[2] - rec1[0]) S_rec2 = (rec2[3] - rec2[1]) * (rec2[2] - rec2[0]) # computing the sum_area sum_area = S_rec1 + S_rec2 # find the each edge of intersect rectangle left_line = max(rec1[0], rec2[0]) right_line = min(rec1[2], rec2[2]) top_line = max(rec1[1], rec2[1]) bottom_line = min(rec1[3], rec2[3]) # judge if there is an intersect if left_line >= right_line or top_line >= bottom_line: return 0 else: intersect = (right_line - left_line) * (bottom_line - top_line) return intersect / (sum_area - intersect) def test_iou(): rect2 = [0, 0, 100, 100] rect1 = (50, 50, 150, 150) iou = compute_iou_numpy(rect1, rect2) print(iou) if __name__ == '__main__': test_iou()
[python][原创]计算IOU和NMS
最新推荐文章于 2024-02-17 16:47:53 发布