Spaces:
Running
Running
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| import logging | |
| import unittest | |
| import torch | |
| from detectron2.config import get_cfg | |
| from detectron2.layers import ShapeSpec | |
| from detectron2.modeling.anchor_generator import DefaultAnchorGenerator, RotatedAnchorGenerator | |
| logger = logging.getLogger(__name__) | |
| class TestAnchorGenerator(unittest.TestCase): | |
| def test_default_anchor_generator(self): | |
| cfg = get_cfg() | |
| cfg.MODEL.ANCHOR_GENERATOR.SIZES = [[32, 64]] | |
| cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS = [[0.25, 1, 4]] | |
| anchor_generator = DefaultAnchorGenerator(cfg, [ShapeSpec(stride=4)]) | |
| # only the last two dimensions of features matter here | |
| num_images = 2 | |
| features = {"stage3": torch.rand(num_images, 96, 1, 2)} | |
| anchors = anchor_generator([features["stage3"]]) | |
| expected_anchor_tensor = torch.tensor( | |
| [ | |
| [-32.0, -8.0, 32.0, 8.0], | |
| [-16.0, -16.0, 16.0, 16.0], | |
| [-8.0, -32.0, 8.0, 32.0], | |
| [-64.0, -16.0, 64.0, 16.0], | |
| [-32.0, -32.0, 32.0, 32.0], | |
| [-16.0, -64.0, 16.0, 64.0], | |
| [-28.0, -8.0, 36.0, 8.0], # -28.0 == -32.0 + STRIDE (4) | |
| [-12.0, -16.0, 20.0, 16.0], | |
| [-4.0, -32.0, 12.0, 32.0], | |
| [-60.0, -16.0, 68.0, 16.0], | |
| [-28.0, -32.0, 36.0, 32.0], | |
| [-12.0, -64.0, 20.0, 64.0], | |
| ] | |
| ) | |
| self.assertTrue(torch.allclose(anchors[0].tensor, expected_anchor_tensor)) | |
| def test_default_anchor_generator_centered(self): | |
| # test explicit args | |
| anchor_generator = DefaultAnchorGenerator( | |
| sizes=[32, 64], aspect_ratios=[0.25, 1, 4], strides=[4] | |
| ) | |
| # only the last two dimensions of features matter here | |
| num_images = 2 | |
| features = {"stage3": torch.rand(num_images, 96, 1, 2)} | |
| expected_anchor_tensor = torch.tensor( | |
| [ | |
| [-30.0, -6.0, 34.0, 10.0], | |
| [-14.0, -14.0, 18.0, 18.0], | |
| [-6.0, -30.0, 10.0, 34.0], | |
| [-62.0, -14.0, 66.0, 18.0], | |
| [-30.0, -30.0, 34.0, 34.0], | |
| [-14.0, -62.0, 18.0, 66.0], | |
| [-26.0, -6.0, 38.0, 10.0], | |
| [-10.0, -14.0, 22.0, 18.0], | |
| [-2.0, -30.0, 14.0, 34.0], | |
| [-58.0, -14.0, 70.0, 18.0], | |
| [-26.0, -30.0, 38.0, 34.0], | |
| [-10.0, -62.0, 22.0, 66.0], | |
| ] | |
| ) | |
| anchors = anchor_generator([features["stage3"]]) | |
| self.assertTrue(torch.allclose(anchors[0].tensor, expected_anchor_tensor)) | |
| anchors = torch.jit.script(anchor_generator)([features["stage3"]]) | |
| self.assertTrue(torch.allclose(anchors[0].tensor, expected_anchor_tensor)) | |
| def test_rrpn_anchor_generator(self): | |
| cfg = get_cfg() | |
| cfg.MODEL.ANCHOR_GENERATOR.SIZES = [[32, 64]] | |
| cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS = [[0.25, 1, 4]] | |
| cfg.MODEL.ANCHOR_GENERATOR.ANGLES = [0, 45] # test single list[float] | |
| anchor_generator = RotatedAnchorGenerator(cfg, [ShapeSpec(stride=4)]) | |
| # only the last two dimensions of features matter here | |
| num_images = 2 | |
| features = {"stage3": torch.rand(num_images, 96, 1, 2)} | |
| anchors = anchor_generator([features["stage3"]]) | |
| expected_anchor_tensor = torch.tensor( | |
| [ | |
| [0.0, 0.0, 64.0, 16.0, 0.0], | |
| [0.0, 0.0, 64.0, 16.0, 45.0], | |
| [0.0, 0.0, 32.0, 32.0, 0.0], | |
| [0.0, 0.0, 32.0, 32.0, 45.0], | |
| [0.0, 0.0, 16.0, 64.0, 0.0], | |
| [0.0, 0.0, 16.0, 64.0, 45.0], | |
| [0.0, 0.0, 128.0, 32.0, 0.0], | |
| [0.0, 0.0, 128.0, 32.0, 45.0], | |
| [0.0, 0.0, 64.0, 64.0, 0.0], | |
| [0.0, 0.0, 64.0, 64.0, 45.0], | |
| [0.0, 0.0, 32.0, 128.0, 0.0], | |
| [0.0, 0.0, 32.0, 128.0, 45.0], | |
| [4.0, 0.0, 64.0, 16.0, 0.0], # 4.0 == 0.0 + STRIDE (4) | |
| [4.0, 0.0, 64.0, 16.0, 45.0], | |
| [4.0, 0.0, 32.0, 32.0, 0.0], | |
| [4.0, 0.0, 32.0, 32.0, 45.0], | |
| [4.0, 0.0, 16.0, 64.0, 0.0], | |
| [4.0, 0.0, 16.0, 64.0, 45.0], | |
| [4.0, 0.0, 128.0, 32.0, 0.0], | |
| [4.0, 0.0, 128.0, 32.0, 45.0], | |
| [4.0, 0.0, 64.0, 64.0, 0.0], | |
| [4.0, 0.0, 64.0, 64.0, 45.0], | |
| [4.0, 0.0, 32.0, 128.0, 0.0], | |
| [4.0, 0.0, 32.0, 128.0, 45.0], | |
| ] | |
| ) | |
| self.assertTrue(torch.allclose(anchors[0].tensor, expected_anchor_tensor)) | |
| if __name__ == "__main__": | |
| unittest.main() | |