|
|
|
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() |
|
|
