27 lines
772 B
Python
27 lines
772 B
Python
from aarecommon.models.models import MLBoxType, MLOutputModel
|
|
|
|
|
|
def test_add_box_generates_unique_keys():
|
|
model = MLOutputModel()
|
|
key1 = model.add_box(MLBoxType.CRYSTAL, (1, 2, 3, 4), 0.8)
|
|
key2 = model.add_box(MLBoxType.CRYSTAL, (5, 6, 7, 8), 0.9)
|
|
|
|
assert key1 == "Crystal"
|
|
assert key2 == "Crystal_2"
|
|
|
|
|
|
def test_get_best_for_class_returns_highest_confidence():
|
|
model = MLOutputModel()
|
|
model.add_box(MLBoxType.PIN, (1, 1, 2, 2), 0.3)
|
|
model.add_box(MLBoxType.PIN, (3, 3, 4, 4), 0.7)
|
|
|
|
best = model.get_best_for_class(MLBoxType.PIN)
|
|
|
|
assert best is not None
|
|
assert best.conf == 0.7
|
|
|
|
|
|
def test_get_best_for_class_returns_none_when_missing():
|
|
model = MLOutputModel()
|
|
assert model.get_best_for_class(MLBoxType.CRYSTAL) is None
|