Spaces:
Configuration error
Configuration error
Download tests/test_localizer.py from ZainabEman/OCR-Based-Spare-Part-Localization: direct link, hf CLI and curl.
- Browser
- Download file 6.73 kB
-
https://huggingface.co/spaces/ZainabEman/OCR-Based-Spare-Part-Localization/resolve/main/tests/test_localizer.py
- Command line
-
hf download hf://spaces/ZainabEman/OCR-Based-Spare-Part-Localization/tests/test_localizer.py
-
curl -L -o test_localizer.py https://huggingface.co/spaces/ZainabEman/OCR-Based-Spare-Part-Localization/resolve/main/tests/test_localizer.py
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" | |
| 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 | |