25 lines
764 B
Python
25 lines
764 B
Python
from aare.common.models import MLOutputModel, MLBoxType
|
|
|
|
|
|
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 |