Spaces:
Running
Running
File size: 1,663 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 |
# Copyright (c) Facebook, Inc. and its affiliates.
import unittest
from typing import List
import torch
from detectron2.config import get_cfg
from detectron2.modeling.matcher import Matcher
class TestMatcher(unittest.TestCase):
def test_scriptability(self):
cfg = get_cfg()
anchor_matcher = Matcher(
cfg.MODEL.RPN.IOU_THRESHOLDS, cfg.MODEL.RPN.IOU_LABELS, allow_low_quality_matches=True
)
match_quality_matrix = torch.tensor(
[[0.15, 0.45, 0.2, 0.6], [0.3, 0.65, 0.05, 0.1], [0.05, 0.4, 0.25, 0.4]]
)
expected_matches = torch.tensor([1, 1, 2, 0])
expected_match_labels = torch.tensor([-1, 1, 0, 1], dtype=torch.int8)
matches, match_labels = anchor_matcher(match_quality_matrix)
self.assertTrue(torch.allclose(matches, expected_matches))
self.assertTrue(torch.allclose(match_labels, expected_match_labels))
# nonzero_tuple must be import explicitly to let jit know what it is.
# https://github.com/pytorch/pytorch/issues/38964
from detectron2.layers import nonzero_tuple # noqa F401
def f(thresholds: List[float], labels: List[int]):
return Matcher(thresholds, labels, allow_low_quality_matches=True)
scripted_anchor_matcher = torch.jit.script(f)(
cfg.MODEL.RPN.IOU_THRESHOLDS, cfg.MODEL.RPN.IOU_LABELS
)
matches, match_labels = scripted_anchor_matcher(match_quality_matrix)
self.assertTrue(torch.allclose(matches, expected_matches))
self.assertTrue(torch.allclose(match_labels, expected_match_labels))
if __name__ == "__main__":
unittest.main()
|