Spaces:
Runtime error
Runtime error
File size: 8,424 Bytes
4a285f6 |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
from __future__ import absolute_import, division, print_function, unicode_literals
import numpy as np
import unittest
import torch
from torchvision import ops
from detectron2.layers import batched_nms, batched_nms_rotated, nms_rotated
def nms_edit_distance(keep1, keep2):
"""
Compare the "keep" result of two nms call.
They are allowed to be different in terms of edit distance
due to floating point precision issues, e.g.,
if a box happen to have an IoU of 0.5 with another box,
one implentation may choose to keep it while another may discard it.
"""
if torch.equal(keep1, keep2):
# they should be equal most of the time
return 0
keep1, keep2 = tuple(keep1.cpu()), tuple(keep2.cpu())
m, n = len(keep1), len(keep2)
# edit distance with DP
f = [np.arange(n + 1), np.arange(n + 1)]
for i in range(m):
cur_row = i % 2
other_row = (i + 1) % 2
f[other_row][0] = i + 1
for j in range(n):
f[other_row][j + 1] = (
f[cur_row][j]
if keep1[i] == keep2[j]
else min(min(f[cur_row][j], f[cur_row][j + 1]), f[other_row][j]) + 1
)
return f[m % 2][n]
class TestNMSRotated(unittest.TestCase):
def reference_horizontal_nms(self, boxes, scores, iou_threshold):
"""
Args:
box_scores (N, 5): boxes in corner-form and probabilities.
(Note here 5 == 4 + 1, i.e., 4-dim horizontal box + 1-dim prob)
iou_threshold: intersection over union threshold.
Returns:
picked: a list of indexes of the kept boxes
"""
picked = []
_, indexes = scores.sort(descending=True)
while len(indexes) > 0:
current = indexes[0]
picked.append(current.item())
if len(indexes) == 1:
break
current_box = boxes[current, :]
indexes = indexes[1:]
rest_boxes = boxes[indexes, :]
iou = ops.box_iou(rest_boxes, current_box.unsqueeze(0)).squeeze(1)
indexes = indexes[iou <= iou_threshold]
return torch.as_tensor(picked)
def _create_tensors(self, N):
boxes = torch.rand(N, 4) * 100
# Note: the implementation of this function in torchvision is:
# boxes[:, 2:] += torch.rand(N, 2) * 100
# but it does not guarantee non-negative widths/heights constraints:
# boxes[:, 2] >= boxes[:, 0] and boxes[:, 3] >= boxes[:, 1]:
boxes[:, 2:] += boxes[:, :2]
scores = torch.rand(N)
return boxes, scores
def test_batched_nms_rotated_0_degree_cpu(self):
N = 2000
num_classes = 50
boxes, scores = self._create_tensors(N)
idxs = torch.randint(0, num_classes, (N,))
rotated_boxes = torch.zeros(N, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
err_msg = "Rotated NMS with 0 degree is incompatible with horizontal NMS for IoU={}"
for iou in [0.2, 0.5, 0.8]:
backup = boxes.clone()
keep_ref = batched_nms(boxes, scores, idxs, iou)
assert torch.allclose(boxes, backup), "boxes modified by batched_nms"
backup = rotated_boxes.clone()
keep = batched_nms_rotated(rotated_boxes, scores, idxs, iou)
assert torch.allclose(
rotated_boxes, backup
), "rotated_boxes modified by batched_nms_rotated"
self.assertLessEqual(nms_edit_distance(keep, keep_ref), 1, err_msg.format(iou))
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
def test_batched_nms_rotated_0_degree_cuda(self):
N = 2000
num_classes = 50
boxes, scores = self._create_tensors(N)
idxs = torch.randint(0, num_classes, (N,))
rotated_boxes = torch.zeros(N, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
err_msg = "Rotated NMS with 0 degree is incompatible with horizontal NMS for IoU={}"
for iou in [0.2, 0.5, 0.8]:
backup = boxes.clone()
keep_ref = batched_nms(boxes.cuda(), scores.cuda(), idxs, iou)
self.assertTrue(torch.allclose(boxes, backup), "boxes modified by batched_nms")
backup = rotated_boxes.clone()
keep = batched_nms_rotated(rotated_boxes.cuda(), scores.cuda(), idxs, iou)
self.assertTrue(
torch.allclose(rotated_boxes, backup),
"rotated_boxes modified by batched_nms_rotated",
)
self.assertLessEqual(nms_edit_distance(keep, keep_ref), 1, err_msg.format(iou))
def test_nms_rotated_0_degree_cpu(self):
N = 1000
boxes, scores = self._create_tensors(N)
rotated_boxes = torch.zeros(N, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
err_msg = "Rotated NMS incompatible between CPU and reference implementation for IoU={}"
for iou in [0.5]:
keep_ref = self.reference_horizontal_nms(boxes, scores, iou)
keep = nms_rotated(rotated_boxes, scores, iou)
self.assertLessEqual(nms_edit_distance(keep, keep_ref), 1, err_msg.format(iou))
def test_nms_rotated_90_degrees_cpu(self):
N = 1000
boxes, scores = self._create_tensors(N)
rotated_boxes = torch.zeros(N, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
# Note for rotated_boxes[:, 2] and rotated_boxes[:, 3]:
# widths and heights are intentionally swapped here for 90 degrees case
# so that the reference horizontal nms could be used
rotated_boxes[:, 2] = boxes[:, 3] - boxes[:, 1]
rotated_boxes[:, 3] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 4] = torch.ones(N) * 90
err_msg = "Rotated NMS incompatible between CPU and reference implementation for IoU={}"
for iou in [0.2, 0.5, 0.8]:
keep_ref = self.reference_horizontal_nms(boxes, scores, iou)
keep = nms_rotated(rotated_boxes, scores, iou)
assert torch.equal(keep, keep_ref), err_msg.format(iou)
def test_nms_rotated_180_degrees_cpu(self):
N = 1000
boxes, scores = self._create_tensors(N)
rotated_boxes = torch.zeros(N, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
rotated_boxes[:, 4] = torch.ones(N) * 180
err_msg = "Rotated NMS incompatible between CPU and reference implementation for IoU={}"
for iou in [0.2, 0.5, 0.8]:
keep_ref = self.reference_horizontal_nms(boxes, scores, iou)
keep = nms_rotated(rotated_boxes, scores, iou)
assert torch.equal(keep, keep_ref), err_msg.format(iou)
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
def test_nms_rotated_0_degree_cuda(self):
N = 1000
boxes, scores = self._create_tensors(N)
rotated_boxes = torch.zeros(N, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
err_msg = "Rotated NMS incompatible between CPU and CUDA for IoU={}"
for iou in [0.2, 0.5, 0.8]:
r_cpu = nms_rotated(rotated_boxes, scores, iou)
r_cuda = nms_rotated(rotated_boxes.cuda(), scores.cuda(), iou)
assert torch.equal(r_cpu, r_cuda.cpu()), err_msg.format(iou)
if __name__ == "__main__":
unittest.main()
|