Wheat_HEAD_Detection_Counting_ComputerVision_Model
/
detectron2
/tests
/modeling
/test_anchor_generator.py
# 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() | |