Spaces:
Configuration error
Configuration error
| """Rendering the red highlight, and leaving the cached diagram untouched.""" | |
| import cv2 | |
| import numpy as np | |
| import pytest | |
| from spareparts import config, index_db, lookup, overlay | |
| pytestmark = pytest.mark.skipif( | |
| not config.DB_PATH.exists(), reason="run `python -m spareparts.build_index` first" | |
| ) | |
| def worm(): | |
| conn = index_db.connect() | |
| try: | |
| part = lookup.find_part(conn, "30642") | |
| finally: | |
| conn.close() | |
| return part, part.occurrences[0] | |
| def _red_pixels(image_bgr): | |
| b, g, r = image_bgr[:, :, 0], image_bgr[:, :, 1], image_bgr[:, :, 2] | |
| return (r > 180) & (g < 80) & (b < 80) | |
| def test_load_diagram_reads_the_cached_png(worm): | |
| _, occurrence = worm | |
| image = overlay.load_diagram(occurrence.png_path) | |
| assert image.shape[:2] == (occurrence.img_h, occurrence.img_w) | |
| def test_load_diagram_raises_on_a_missing_file(tmp_path): | |
| with pytest.raises(FileNotFoundError): | |
| overlay.load_diagram(tmp_path / "absent.png") | |
| def test_annotation_does_not_modify_the_cached_image_on_disk(worm): | |
| part, occurrence = worm | |
| before = occurrence.png_path.read_bytes() | |
| image = overlay.load_diagram(occurrence.png_path) | |
| overlay.annotate(image, occurrence.highlights, part.part_no) | |
| assert occurrence.png_path.read_bytes() == before | |
| def test_annotation_does_not_modify_the_array_passed_in(worm): | |
| part, occurrence = worm | |
| image = overlay.load_diagram(occurrence.png_path) | |
| original = image.copy() | |
| annotated = overlay.annotate(image, occurrence.highlights, part.part_no) | |
| assert np.array_equal(image, original), "the source array was drawn on" | |
| assert not np.array_equal(annotated, original), "nothing was drawn" | |
| def test_red_rectangle_lands_on_the_stored_box(worm): | |
| part, occurrence = worm | |
| x0, y0, x1, y1 = occurrence.highlights[0].component_box | |
| image = overlay.load_diagram(occurrence.png_path) | |
| annotated = overlay.annotate( | |
| image, occurrence.highlights, part.part_no, show_callout_box=False | |
| ) | |
| red = _red_pixels(annotated) | |
| # Red appears along each edge of the stored box... | |
| margin = 8 | |
| assert red[max(0, y0 - margin) : y0 + margin, x0:x1].any(), "no top edge" | |
| assert red[y1 - margin : y1 + margin, x0:x1].any(), "no bottom edge" | |
| assert red[y0:y1, max(0, x0 - margin) : x0 + margin].any(), "no left edge" | |
| assert red[y0:y1, x1 - margin : x1 + margin].any(), "no right edge" | |
| # ...and nowhere far away from it, apart from the caption above the box. | |
| outside = red.copy() | |
| outside[max(0, y0 - 140) : y1 + margin, max(0, x0 - margin) : x1 + margin] = False | |
| assert outside.sum() < 50, "red drawn well outside the stored box" | |
| def test_caption_shows_the_part_number(worm): | |
| """The number is drawn on a white plate near the box, not left implicit.""" | |
| part, occurrence = worm | |
| x0, y0, _, _ = occurrence.highlights[0].component_box | |
| image = overlay.load_diagram(occurrence.png_path) | |
| annotated = overlay.annotate(image, occurrence.highlights, part.part_no) | |
| plate = annotated[max(0, y0 - 130) : y0, max(0, x0 - 10) : x0 + 400] | |
| # A white plate with red text on it. | |
| assert (plate == 255).all(axis=2).sum() > 500 | |
| assert _red_pixels(plate).any() | |
| def test_every_instance_is_boxed_and_numbered(): | |
| """74506 is drawn four times on page 34; all four must be highlighted.""" | |
| conn = index_db.connect() | |
| try: | |
| part = lookup.find_part(conn, "74506") | |
| finally: | |
| conn.close() | |
| occurrence = next(o for o in part.occurrences if o.pdf_page == 34) | |
| image = overlay.load_diagram(occurrence.png_path) | |
| annotated = overlay.annotate( | |
| image, occurrence.highlights, part.part_no, show_callout_box=False | |
| ) | |
| red = _red_pixels(annotated) | |
| for highlight in occurrence.highlights: | |
| x0, y0, x1, y1 = highlight.component_box | |
| assert red[y0 : y1 + 1, x0 : x1 + 1].any(), f"instance at {highlight.component_box}" | |
| def test_missing_component_box_is_skipped_not_crashed(worm): | |
| part, occurrence = worm | |
| blank = lookup.Highlight( | |
| instance=0, label_box=(10, 10, 100, 40), component_box=None, | |
| method="none", confidence=0.0, note="", | |
| ) | |
| image = overlay.load_diagram(occurrence.png_path) | |
| annotated = overlay.annotate(image, [blank], part.part_no) | |
| assert annotated.shape == image.shape | |
| def test_box_outside_the_image_is_clipped(worm): | |
| part, occurrence = worm | |
| huge = lookup.Highlight( | |
| instance=0, label_box=(0, 0, 50, 20), | |
| component_box=(-500, -500, occurrence.img_w + 999, occurrence.img_h + 999), | |
| method="leader", confidence=1.0, note="", | |
| ) | |
| image = overlay.load_diagram(occurrence.png_path) | |
| annotated = overlay.annotate(image, [huge], part.part_no) | |
| assert annotated.shape == image.shape | |
| def test_encode_png_round_trips(worm): | |
| part, occurrence = worm | |
| image = overlay.load_diagram(occurrence.png_path) | |
| data = overlay.encode_png(overlay.annotate(image, occurrence.highlights, part.part_no)) | |
| assert data[:8] == b"\x89PNG\r\n\x1a\n" | |
| decoded = cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR) | |
| assert decoded.shape == image.shape | |