def test_iou_0_dim_cuda(self):
boxes1 = torch.rand(0, 5, dtype=torch.float32)
boxes2 = torch.rand(10, 5, dtype=torch.float32)
expected_ious = torch.zeros(0, 10, dtype=torch.float32)
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
self.assertTrue(torch.allclose(ious_cuda.cpu(), expected_ious))
boxes1 = torch.rand(10, 5, dtype=torch.float32)
boxes2 = torch.rand(0, 5, dtype=torch.float32)
expected_ious = torch.zeros(10, 0, dtype=torch.float32)
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
self.assertTrue(torch.allclose(ious_cuda.cpu(), expected_ious))
def test_iou_half_overlap_cuda(self):
boxes1 = torch.tensor([[0.5, 0.5, 1.0, 1.0, 0.0]], dtype=torch.float32)
boxes2 = torch.tensor([[0.25, 0.5, 0.5, 1.0, 0.0]],
dtype=torch.float32)
expected_ious = torch.tensor([[0.5]], dtype=torch.float32)
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
self.assertTrue(torch.allclose(ious_cuda.cpu(), expected_ious))
def test_iou_perpendicular_cuda(self):
boxes1 = torch.tensor([[5, 5, 10.0, 6, 55]], dtype=torch.float32)
boxes2 = torch.tensor([[5, 5, 10.0, 6, -35]], dtype=torch.float32)
iou = (6.0 * 6.0) / (6.0 * 6.0 + 4.0 * 6.0 + 4.0 * 6.0)
expected_ious = torch.tensor([[iou]], dtype=torch.float32)
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
self.assertTrue(torch.allclose(ious_cuda.cpu(), expected_ious))
def test_iou_precision_cuda(self):
boxes1 = torch.tensor([[565, 565, 10, 10, 0]], dtype=torch.float32)
boxes2 = torch.tensor([[565, 565, 10, 8.3, 0]], dtype=torch.float32)
iou = 8.3 / 10.0
expected_ious = torch.tensor([[iou]], dtype=torch.float32)
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
self.assertTrue(torch.allclose(ious_cuda.cpu(), expected_ious))
def test_iou_large_close_boxes_cuda(self):
boxes1 = torch.tensor(
[[299.500000, 417.370422, 600.000000, 364.259186, 27.1828]],
dtype=torch.float32)
boxes2 = torch.tensor(
[[299.500000, 417.370422, 600.000000, 364.259155, 27.1828]],
dtype=torch.float32)
iou = 364.259155 / 364.259186
expected_ious = torch.tensor([[iou]], dtype=torch.float32)
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
assert torch.allclose(ious_cuda.cpu(), expected_ious)
def test_iou_large_close_boxes_cpu(self):
boxes1 = torch.tensor(
[[299.500000, 417.370422, 600.000000, 364.259186, 27.1828]],
dtype=torch.float32)
boxes2 = torch.tensor(
[[299.500000, 417.370422, 600.000000, 364.259155, 27.1828]],
dtype=torch.float32)
iou = 364.259155 / 364.259186
expected_ious = torch.tensor([[iou]], dtype=torch.float32)
ious = pairwise_iou_rotated(boxes1, boxes2)
self.assertTrue(torch.allclose(ious, expected_ious))
def test_iou_45_degrees_cuda(self):
boxes1 = torch.tensor(
[
[1, 1, math.sqrt(2), math.sqrt(2), 45],
[1, 1, 2 * math.sqrt(2), 2 * math.sqrt(2), -45],
],
dtype=torch.float32,
)
boxes2 = torch.tensor([[1, 1, 2, 2, 0]], dtype=torch.float32)
expected_ious = torch.tensor([[0.5], [0.5]], dtype=torch.float32)
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
self.assertTrue(torch.allclose(ious_cuda.cpu(), expected_ious))
def test_iou_45_degrees_cpu(self):
boxes1 = torch.tensor(
[
[1, 1, math.sqrt(2), math.sqrt(2), 45],
[1, 1, 2 * math.sqrt(2), 2 * math.sqrt(2), -45],
],
dtype=torch.float32,
)
boxes2 = torch.tensor([[1, 1, 2, 2, 0]], dtype=torch.float32)
expected_ious = torch.tensor([[0.5], [0.5]], dtype=torch.float32)
ious = pairwise_iou_rotated(boxes1, boxes2)
assert torch.allclose(ious, expected_ious)
def pairwise_iou(boxes1: RotatedBoxes, boxes2: RotatedBoxes) -> None:
"""
Given two lists of rotated boxes of size N and M,
compute the IoU (intersection over union)
between __all__ N x M pairs of boxes.
The box order must be (x_center, y_center, width, height, angle).
Args:
boxes1, boxes2 (RotatedBoxes):
two `RotatedBoxes`. Contains N & M rotated boxes, respectively.
Returns:
Tensor: IoU, sized [N,M].
"""
return pairwise_iou_rotated(boxes1.tensor, boxes2.tensor)
def test_iou_many_boxes_cuda(self):
num_boxes1 = 100
num_boxes2 = 200
boxes1 = torch.stack([
torch.tensor([5 + 20 * i, 5 + 20 * i, 10, 10, 0],
dtype=torch.float32) for i in range(num_boxes1)
])
boxes2 = torch.stack([
torch.tensor(
[5 + 20 * i, 5 + 20 * i, 10, 1 + 9 * i / num_boxes2, 0],
dtype=torch.float32) for i in range(num_boxes2)
])
expected_ious = torch.zeros(num_boxes1,
num_boxes2,
dtype=torch.float32)
for i in range(min(num_boxes1, num_boxes2)):
expected_ious[i][i] = (1 + 9 * i / num_boxes2) / 10.0
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
self.assertTrue(torch.allclose(ious_cuda.cpu(), expected_ious))
def test_iou_0_degree_cuda(self):
boxes1 = torch.tensor(
[[0.5, 0.5, 1.0, 1.0, 0.0], [0.5, 0.5, 1.0, 1.0, 0.0]],
dtype=torch.float32)
boxes2 = torch.tensor(
[
[0.5, 0.5, 1.0, 1.0, 0.0],
[0.25, 0.5, 0.5, 1.0, 0.0],
[0.5, 0.25, 1.0, 0.5, 0.0],
[0.25, 0.25, 0.5, 0.5, 0.0],
[0.75, 0.75, 0.5, 0.5, 0.0],
[1.0, 1.0, 1.0, 1.0, 0.0],
],
dtype=torch.float32,
)
expected_ious = torch.tensor(
[
[1.0, 0.5, 0.5, 0.25, 0.25, 0.25 / (2 - 0.25)],
[1.0, 0.5, 0.5, 0.25, 0.25, 0.25 / (2 - 0.25)],
],
dtype=torch.float32,
)
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
self.assertTrue(torch.allclose(ious_cuda.cpu(), expected_ious))
def bench():
for _ in range(n):
pairwise_iou_rotated(b1, b2)
if dev.type == "cuda":
torch.cuda.synchronize()
def test_iou_too_many_boxes_cuda(self):
s1, s2 = 5, 1289035
boxes1 = torch.zeros(s1, 5)
boxes2 = torch.zeros(s2, 5)
ious_cuda = pairwise_iou_rotated(boxes1.cuda(), boxes2.cuda())
self.assertTupleEqual(tuple(ious_cuda.shape), (s1, s2))