Spaces:
Running
Running
File size: 2,003 Bytes
6f6830f |
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 |
# Copyright (c) Facebook, Inc. and its affiliates.
import unittest
import torch
from detectron2.structures.masks import BitMasks, PolygonMasks, polygons_to_bitmask
class TestBitMask(unittest.TestCase):
def test_get_bounding_box(self):
masks = torch.tensor(
[
[
[False, False, False, True],
[False, False, True, True],
[False, True, True, False],
[False, True, True, False],
],
[
[False, False, False, False],
[False, False, True, False],
[False, True, True, False],
[False, True, True, False],
],
torch.zeros(4, 4),
]
)
bitmask = BitMasks(masks)
box_true = torch.tensor([[1, 0, 4, 4], [1, 1, 3, 4], [0, 0, 0, 0]], dtype=torch.float32)
box = bitmask.get_bounding_boxes()
self.assertTrue(torch.all(box.tensor == box_true).item())
for box in box_true:
poly = box[[0, 1, 2, 1, 2, 3, 0, 3]].numpy()
mask = polygons_to_bitmask([poly], 4, 4)
reconstruct_box = BitMasks(mask[None, :, :]).get_bounding_boxes()[0].tensor
self.assertTrue(torch.all(box == reconstruct_box).item())
reconstruct_box = PolygonMasks([[poly]]).get_bounding_boxes()[0].tensor
self.assertTrue(torch.all(box == reconstruct_box).item())
def test_from_empty_polygons(self):
masks = BitMasks.from_polygon_masks([], 100, 100)
self.assertEqual(masks.tensor.shape, (0, 100, 100))
def test_getitem(self):
masks = BitMasks(torch.ones(3, 10, 10))
self.assertEqual(masks[1].tensor.shape, (1, 10, 10))
self.assertEqual(masks[1:3].tensor.shape, (2, 10, 10))
self.assertEqual(masks[torch.tensor([True, False, False])].tensor.shape, (1, 10, 10))
if __name__ == "__main__":
unittest.main()
|