41 lines
1.4 KiB
Python
41 lines
1.4 KiB
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)
|
|
|
|
def test_best_by_class_keeps_highest_confidence_per_class():
|
|
results = [
|
|
{"class": 1, "confidence": 0.4, "box": {"x1": 0, "y1": 0, "x2": 10, "y2": 10}},
|
|
{"class": 1, "confidence": 0.8, "box": {"x1": 1, "y1": 1, "x2": 11, "y2": 11}},
|
|
{"class": 2, "confidence": 0.6, "box": {"x1": 2, "y1": 2, "x2": 12, "y2": 12}},
|
|
]
|
|
|
|
out = MlBox._best_by_class(results)
|
|
|
|
assert out is not None
|
|
assert out.get_best_for_class(MLBoxType.Pin).conf == 0.8
|
|
assert out.get_best_for_class(MLBoxType.Crystal).conf == 0.6
|