46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
from scripts.audit_yolo_label_relationships import Box, audit_boxes, classify_pair
|
|
|
|
|
|
DEFAULTS = {
|
|
"near_duplicate_iou": 0.9,
|
|
"containment_threshold": 0.98,
|
|
"max_nested_area_ratio": 4.0,
|
|
}
|
|
|
|
|
|
def test_classifies_exact_near_and_possible_nested_relationships() -> None:
|
|
exact = Box(0, 0.5, 0.5, 0.2, 0.2)
|
|
near = Box(0, 0.501, 0.5, 0.2, 0.2)
|
|
outer = Box(0, 0.5, 0.5, 0.18, 0.18)
|
|
inner = Box(0, 0.5, 0.5, 0.1, 0.1)
|
|
|
|
assert classify_pair(exact, exact, **DEFAULTS) == ("exact_duplicate", 1.0)
|
|
assert classify_pair(exact, near, **DEFAULTS)[0] == "near_duplicate"
|
|
assert classify_pair(outer, inner, **DEFAULTS) == ("possible_nested", 1.0)
|
|
|
|
|
|
def test_ignores_other_classes_and_non_overlapping_boxes() -> None:
|
|
first = Box(0, 0.2, 0.2, 0.1, 0.1)
|
|
other_class = Box(1, 0.2, 0.2, 0.1, 0.1)
|
|
distant = Box(0, 0.8, 0.8, 0.1, 0.1)
|
|
|
|
assert classify_pair(first, other_class, **DEFAULTS) is None
|
|
assert classify_pair(first, distant, **DEFAULTS) is None
|
|
|
|
|
|
def test_audit_boxes_returns_pair_indices_and_coordinates() -> None:
|
|
first = Box(0, 0.5, 0.5, 0.2, 0.2)
|
|
second = Box(0, 0.5, 0.5, 0.2, 0.2)
|
|
relationships = audit_boxes([first, second], **DEFAULTS)
|
|
|
|
assert relationships == [
|
|
{
|
|
"relationship": "exact_duplicate",
|
|
"score": 1.0,
|
|
"first_index": 0,
|
|
"second_index": 1,
|
|
"first_box": [0, 0.5, 0.5, 0.2, 0.2],
|
|
"second_box": [0, 0.5, 0.5, 0.2, 0.2],
|
|
}
|
|
]
|