File size: 1,509 Bytes
9b92c75 | 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 | import torch
from objectmodel_v1.matching import hungarian_match, hungarian_match_layers
def test_layer_matcher_matches_individual_calls():
torch.manual_seed(7)
outputs = [
{
"pred_logits": torch.randn(3, 12, 5),
"pred_boxes": torch.rand(3, 12, 4),
}
for _ in range(3)
]
targets = [
{"labels": torch.tensor([1, 3]), "boxes": torch.rand(2, 4)},
{"labels": torch.empty(0, dtype=torch.long), "boxes": torch.empty(0, 4)},
{"labels": torch.tensor([0, 2, 4]), "boxes": torch.rand(3, 4)},
]
expected = [hungarian_match(output, targets) for output in outputs]
actual = hungarian_match_layers(outputs, targets)
for expected_layer, actual_layer in zip(expected, actual, strict=True):
for expected_match, actual_match in zip(expected_layer, actual_layer, strict=True):
assert torch.equal(expected_match[0], actual_match[0])
assert torch.equal(expected_match[1], actual_match[1])
def test_layer_matcher_handles_all_empty_targets():
outputs = [{"pred_logits": torch.randn(2, 4, 3), "pred_boxes": torch.rand(2, 4, 4)}]
targets = [
{"labels": torch.empty(0, dtype=torch.long), "boxes": torch.empty(0, 4)},
{"labels": torch.empty(0, dtype=torch.long), "boxes": torch.empty(0, 4)},
]
matches = hungarian_match_layers(outputs, targets)
assert len(matches) == 1
assert all(rows.numel() == 0 and cols.numel() == 0 for rows, cols in matches[0])
|