File size: 3,647 Bytes
0f90f73 |
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 |
# Copyright (c) Facebook, Inc. and its affiliates.
import logging
import unittest
import torch
from detectron2.modeling.box_regression import (
Box2BoxTransform,
Box2BoxTransformLinear,
Box2BoxTransformRotated,
)
from detectron2.utils.testing import random_boxes
logger = logging.getLogger(__name__)
class TestBox2BoxTransform(unittest.TestCase):
def test_reconstruction(self):
weights = (5, 5, 10, 10)
b2b_tfm = Box2BoxTransform(weights=weights)
src_boxes = random_boxes(10)
dst_boxes = random_boxes(10)
devices = [torch.device("cpu")]
if torch.cuda.is_available():
devices.append(torch.device("cuda"))
for device in devices:
src_boxes = src_boxes.to(device=device)
dst_boxes = dst_boxes.to(device=device)
deltas = b2b_tfm.get_deltas(src_boxes, dst_boxes)
dst_boxes_reconstructed = b2b_tfm.apply_deltas(deltas, src_boxes)
self.assertTrue(torch.allclose(dst_boxes, dst_boxes_reconstructed))
def test_apply_deltas_tracing(self):
weights = (5, 5, 10, 10)
b2b_tfm = Box2BoxTransform(weights=weights)
with torch.no_grad():
func = torch.jit.trace(b2b_tfm.apply_deltas, (torch.randn(10, 20), torch.randn(10, 4)))
o = func(torch.randn(10, 20), torch.randn(10, 4))
self.assertEqual(o.shape, (10, 20))
o = func(torch.randn(5, 20), torch.randn(5, 4))
self.assertEqual(o.shape, (5, 20))
def random_rotated_boxes(mean_box, std_length, std_angle, N):
return torch.cat(
[torch.rand(N, 4) * std_length, torch.rand(N, 1) * std_angle], dim=1
) + torch.tensor(mean_box, dtype=torch.float)
class TestBox2BoxTransformRotated(unittest.TestCase):
def test_reconstruction(self):
weights = (5, 5, 10, 10, 1)
b2b_transform = Box2BoxTransformRotated(weights=weights)
src_boxes = random_rotated_boxes([10, 10, 20, 20, -30], 5, 60.0, 10)
dst_boxes = random_rotated_boxes([10, 10, 20, 20, -30], 5, 60.0, 10)
devices = [torch.device("cpu")]
if torch.cuda.is_available():
devices.append(torch.device("cuda"))
for device in devices:
src_boxes = src_boxes.to(device=device)
dst_boxes = dst_boxes.to(device=device)
deltas = b2b_transform.get_deltas(src_boxes, dst_boxes)
dst_boxes_reconstructed = b2b_transform.apply_deltas(deltas, src_boxes)
assert torch.allclose(dst_boxes[:, :4], dst_boxes_reconstructed[:, :4], atol=1e-5)
# angle difference has to be normalized
assert torch.allclose(
(dst_boxes[:, 4] - dst_boxes_reconstructed[:, 4] + 180.0) % 360.0 - 180.0,
torch.zeros_like(dst_boxes[:, 4]),
atol=1e-4,
)
class TestBox2BoxTransformLinear(unittest.TestCase):
def test_reconstruction(self):
b2b_tfm = Box2BoxTransformLinear()
src_boxes = random_boxes(10)
dst_boxes = torch.tensor([0, 0, 101, 101] * 10).reshape(10, 4).float()
devices = [torch.device("cpu")]
if torch.cuda.is_available():
devices.append(torch.device("cuda"))
for device in devices:
src_boxes = src_boxes.to(device=device)
dst_boxes = dst_boxes.to(device=device)
deltas = b2b_tfm.get_deltas(src_boxes, dst_boxes)
dst_boxes_reconstructed = b2b_tfm.apply_deltas(deltas, src_boxes)
self.assertTrue(torch.allclose(dst_boxes, dst_boxes_reconstructed, atol=1e-3))
if __name__ == "__main__":
unittest.main()
|