ZainabEman's picture
Upload 152 files
546c207 verified
Raw History Blame Contribute Delete
6.73 kB
"""Leader-line tracing and component resolution, on synthetic geometry.
The canvas is the real diagram size (1689 x ~2350) and the strokes are the real
weights, because the localizer's thresholds are ratios of the page: testing at a
smaller scale would exercise different behaviour than production sees.
The crossing case is the failure mode PSD.md section 4.9 records: a leader line
passing through a large structure on its way to a small part.
"""
import cv2
import numpy as np
import pytest
from spareparts import localizer
from spareparts.matcher import Label
W, H = 1689, 2352
INK = (0, 0, 0)
STROKE = 3
def canvas():
return np.full((H, W, 3), 255, np.uint8)
def label_at(part_no, x0, y0):
"""A callout box the size OCR actually returns for a 5-digit number."""
return Label(
part_no=part_no, raw_text=part_no, confidence=1.0, ocr_confidence=1.0,
x0=x0, y0=y0, x1=x0 + 95, y1=y0 + 33,
match_kind="exact", has_asterisk=False, in_page_table=True,
)
def draw_digits(img, label):
"""Digit-like strokes wholly inside the callout box."""
for i in range(5):
x = label.x0 + 8 + i * 18
cv2.line(img, (x, label.y0 + 6), (x, label.y1 - 6), INK, STROKE)
def localize(img, labels, overrides=None):
return localizer.localize_page(
img, labels, list(range(len(labels))), overrides, pdf_page=99
)
def covers(box, target, slack=40):
"""The prediction encloses the target without being wildly larger."""
x0, y0, x1, y1 = box
tx0, ty0, tx1, ty1 = target
return (
x0 <= tx0 + slack and y0 <= ty0 + slack
and x1 >= tx1 - slack and y1 >= ty1 - slack
and (x1 - x0) < (tx1 - tx0) * 2.5
and (y1 - y0) < (ty1 - ty0) * 2.5
)
def test_leader_reaches_its_target_and_ignores_a_nearer_decoy():
img = canvas()
label = label_at("30642", 1100, 1500)
draw_digits(img, label)
# A decoy part, closer to the callout than the real target is.
cv2.rectangle(img, (900, 1300), (1010, 1420), INK, STROKE)
# The real target, reached by a leader running up and to the left.
cv2.circle(img, (400, 600), 110, INK, STROKE)
cv2.line(img, (1119, 1482), (469, 685), INK, STROKE)
result = localize(img, [label])[0]
assert result.method == "leader"
assert covers(result.box, (290, 490, 510, 710))
# Nowhere near the decoy.
assert result.box[0] < 700
def test_leader_crossing_a_large_structure_finds_the_small_target_beyond_it():
"""PSD 4.9: on page 34 this collapsed onto the gearbox housing."""
img = canvas()
label = label_at("74506", 1250, 1180)
draw_digits(img, label)
housing = (500, 900, 1150, 1500) # ~10% of the page, as on page 34
cv2.rectangle(img, housing[:2], housing[2:], INK, STROKE)
cv2.circle(img, (330, 1196), 60, INK, STROKE) # the actual target
cv2.line(img, (1245, 1196), (390, 1196), INK, STROKE) # crosses the housing
result = localize(img, [label])[0]
assert result.method == "leader"
assert covers(result.box, (270, 1136, 390, 1256))
# The housing must not have been asserted as the answer.
assert result.box[2] < housing[0]
def test_callout_with_no_leader_line_falls_back_conservatively():
img = canvas()
label = label_at("73624", 700, 700)
draw_digits(img, label)
cv2.circle(img, (750, 950), 90, INK, STROKE)
result = localize(img, [label])[0]
assert result.method == "nearest"
assert result.confidence < 0.5 # a weaker claim, and it says so
assert result.note
def test_oversized_target_is_capped_rather_than_asserted():
"""A leader ending on a page-spanning blob must not box the whole blob."""
img = canvas()
label = label_at("30642", 1400, 2100)
draw_digits(img, label)
for y in range(100, 2000, 7): # a dense blob over most of the page
cv2.line(img, (100, y), (1500, y), INK, 2)
cv2.line(img, (1411, 2088), (1250, 1960), INK, STROKE)
result = localize(img, [label])[0]
assert result.method == "capped"
x0, y0, x1, y1 = result.box
assert (x1 - x0) <= 2 * localizer.FALLBACK_BOX_HALF + 2
assert result.confidence < 0.5
assert result.note # and it says why it was capped
def test_override_wins_and_is_labelled_as_such():
img = canvas()
label = label_at("30642", 1100, 1500)
draw_digits(img, label)
cv2.circle(img, (400, 600), 110, INK, STROKE)
cv2.line(img, (1119, 1482), (469, 685), INK, STROKE)
overrides = {"99:30642:0": {"box": [10, 20, 300, 400], "reason": "ambiguous"}}
result = localize(img, [label], overrides)[0]
assert result.method == "override"
assert result.box == (10, 20, 300, 400)
assert result.note == "ambiguous"
assert result.confidence == 1.0
def test_boxed_callout_is_detected_as_an_assembly():
img = canvas()
label = label_at("31269", 1100, 1500)
draw_digits(img, label)
# The rectangle the catalogue draws around 'kpl.' callouts, just inside the
# OCR box as it is in the real diagrams.
cv2.rectangle(img, (label.x0 + 3, label.y0 + 3), (label.x1 - 3, label.y1 - 3), INK, 2)
cv2.circle(img, (400, 600), 110, INK, STROKE)
cv2.line(img, (1119, 1478), (469, 685), INK, STROKE)
geometry = localizer._PageGeometry(img, [label])
assert 0 in geometry.boxed_callouts
result = localize(img, [label])[0]
assert result.method == "assembly"
@pytest.mark.parametrize("dashed", [False, True])
def test_solid_and_dashed_leaders_both_trace(dashed):
img = canvas()
label = label_at("73158", 1200, 1180)
draw_digits(img, label)
cv2.circle(img, (400, 1196), 90, INK, STROKE)
if dashed:
for x in range(500, 1195, 26): # page 40 uses dashed leaders
cv2.line(img, (x, 1196), (x + 16, 1196), INK, STROKE)
else:
cv2.line(img, (1195, 1196), (500, 1196), INK, STROKE)
result = localize(img, [label])[0]
assert result.method == "leader"
assert covers(result.box, (310, 1106, 490, 1286))
def test_every_result_reports_the_tier_that_produced_it():
img = canvas()
label = label_at("30642", 1100, 1500)
draw_digits(img, label)
cv2.circle(img, (400, 600), 110, INK, STROKE)
cv2.line(img, (1119, 1482), (469, 685), INK, STROKE)
result = localize(img, [label])[0]
assert result.method in {
"leader", "assembly", "capped", "nearest", "override", "none"
}
assert 0.0 <= result.confidence <= 1.0
assert result.anchor is not None
assert len(result.leader) > 10