[目标检测]IoU计算

发表于 更新于 分类于
本文字数： 2.1k 阅读时长 ≈ 4 分钟

之前学习IoU的概念并且实现了预测框和对应真值边界框之间的计算 - [目标检测]IoU

不过预设的条件是每个真值边界框仅和单个预测边界框进行比对，参考box_utils.py计算每个真值边界框和每个预测边界框的IoU

实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import numpy as np
import torch
import cv2


def area_of(left_top, right_bottom) -> torch.Tensor:
    """Compute the areas of rectangles given two corners.

    Args:
        left_top (N, 2): left top corner.
        right_bottom (N, 2): right bottom corner.

    Returns:
        area (N): return the area.
    """
    wh = torch.clamp(right_bottom - left_top, min=0.0)
    return wh[..., 0] * wh[..., 1]


def iou_of(pred_boxes, target_boxes, eps=1e-5):
    """Return intersection-over-union (Jaccard index) of boxes.

    Args:
        pred_boxes (N1, 4): predicted boxes.
        target_boxes (N2, 4): ground truth boxes.
        eps: a small number to avoid 0 as denominator.
    Returns:
        iou (N): IoU values.
    """
    # (N1, 4) -> (1, N1, 4)
    pred_boxes = pred_boxes.unsqueeze(0)
    # (N2, 4) -> (N2, 1, 4)
    target_boxes = target_boxes.unsqueeze(1)

    # (N2, N1, 2)
    overlap_left_top = torch.max(pred_boxes[..., :2], target_boxes[..., :2])
    # (N2, N1, 2)
    overlap_right_bottom = torch.min(pred_boxes[..., 2:], target_boxes[..., 2:])

    # (N2, N1)
    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
    # (1, N1)
    area0 = area_of(pred_boxes[..., :2], pred_boxes[..., 2:])
    # (N2, 1)
    area1 = area_of(target_boxes[..., :2], target_boxes[..., 2:])
    return overlap_area / (area0 + area1 - overlap_area + eps)

测试

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
if __name__ == '__main__':
    pred_boxes = np.array([[30, 280, 200, 300], [20, 300, 100, 360]])
    target_boxes = np.array([[10, 10, 50, 50], [40, 270, 100, 380], [450, 300, 500, 500]])

    # ious大小为[N2, N1]
    # N1表示pred_boxes的个数
    # N2表示target_boxes的个数
    ious = iou_of(torch.from_numpy(pred_boxes), torch.from_numpy(target_boxes))

    print('ious: ', ious)

    img = np.ones((650, 650, 3)) * 255
    for item in target_boxes:
        xmin, ymin, xmax, ymax = item
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), thickness=2)
    for item in pred_boxes:
        xmin, ymin, xmax, ymax = item
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), thickness=1)

    cv2.imshow('img', img)
    cv2.waitKey(0)

    cv2.imwrite('iou.png', img)
###################### 输出每个GT和预测边界框的IoU
ious:  tensor([[0.0000, 0.0000],
        [0.1364, 0.4615],
        [0.0000, 0.0000]])