python实现iou计算

两个boxes之间的iou计算：

import numpy as np

def iou(box1, box2):
    x1, y1, x2, y2 = box1 
    w1, h1, w2, h2 = box2 
    left_max = max(x1, w1)
    right_min = min(x2, w2)
    top_max = max(y1, h1)
    bot_min = min(y2, h2)
    if left_max >= right_min or top_max >= bot_min:
        return 0 
    else:
        intersection = (right_min - left_max) * (bot_min - top_max)
        union = (x2-x1) * (y2 - y1) + (w2 - w1) * (h2 - h1) - intersection
        return intersection / union

b1 = np.array([0, 0, 200, 200])
b2 = np.array([100, 100, 500, 500])
out = iou(b1, b2)   # out = 0.0526

　　

两组boxes的两两box之间的iou计算，下面用的pytorch的张量实现(源码来自torchvision)：

import torch
def box_iou(boxes1, boxes2):
    '''
    boxes1和boxes2的shape都是: (n_box, 4). 4是指(x1, y1, x2, y2) coordinates.
    '''
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2] 这里使用了广播机制
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]

    union = area1[:, None] + area2 - inter

    iou = inter / union

    return iou, union

def box_area(boxes):
    """
    Computes the area of a set of bounding boxes, which are specified by their
    (x1, y1, x2, y2) coordinates.
    """
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])