[python][原创]计算IOU和NMS

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()