27 lines
847 B
Python
27 lines
847 B
Python
from aare.daq.mlbox import MlBox
|
|
from aare.common.models import MLOutputModel, MLBoxType
|
|
|
|
|
|
def test_preferred_class_prefers_loop_over_pin_when_margin_not_exceeded():
|
|
boxes = MLOutputModel()
|
|
boxes.add_box(MLBoxType.Pin, (0, 0, 10, 10), 0.8)
|
|
boxes.add_box(MLBoxType.Loop_face, (1, 1, 9, 9), 0.75)
|
|
|
|
best = MlBox.get_preferred_class_box_with_confidence_threshold(
|
|
boxes,
|
|
loop_preference_margin=0.1,
|
|
)
|
|
|
|
assert best is not None
|
|
assert best.cls == MLBoxType.Loop_face
|
|
|
|
def test_prediction_score_empty():
|
|
assert MlBox._prediction_score(None) == (0, 0.0)
|
|
|
|
|
|
def test_prediction_score_uses_count_and_max_confidence():
|
|
model = MLOutputModel()
|
|
model.add_box(MLBoxType.Pin, (0, 0, 10, 10), 0.4)
|
|
model.add_box(MLBoxType.Crystal, (1, 1, 9, 9), 0.8)
|
|
|
|
assert MlBox._prediction_score(model) == (2, 0.8) |