[目标检测]IoU计算

之前学习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]])

坚持原创技术分享,您的支持将鼓励我继续创作!