Spaces:
Sleeping
Sleeping
File size: 4,626 Bytes
d9106f7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """
Unit tests for OCRProcessor.
"""
import cv2
import numpy as np
import pytest
from processors.ocr_processor import OCRProcessor
@pytest.fixture
def processor():
return OCRProcessor()
@pytest.fixture
def text_drawing():
"""Large, clear image with well-spaced text for reliable OCR."""
img = np.ones((900, 1200), dtype=np.uint8) * 255
font = cv2.FONT_HERSHEY_SIMPLEX
cv2.putText(img, "Living Room", (100, 200), font, 2.0, 0, 4)
cv2.putText(img, "3640", (100, 450), font, 2.0, 0, 4)
cv2.putText(img, "Bedroom", (100, 700), font, 2.0, 0, 4)
return img
class TestExtractText:
def test_returns_dict_with_expected_keys(self, processor, synthetic_drawing):
result = processor.extract_text(synthetic_drawing)
for key in ("full_text", "element_count", "elements", "classified"):
assert key in result, f"Missing key: {key}"
def test_classified_has_correct_categories(self, processor, synthetic_drawing):
result = processor.extract_text(synthetic_drawing)
for cat in ("dimensions", "room_names", "numbers", "labels"):
assert cat in result["classified"]
def test_element_count_matches_elements_list(self, processor, text_drawing):
result = processor.extract_text(text_drawing)
assert result["element_count"] == len(result["elements"])
def test_each_element_has_required_fields(self, processor, text_drawing):
result = processor.extract_text(text_drawing)
for el in result["elements"]:
for field in ("text", "confidence", "x", "y", "width", "height"):
assert field in el, f"Missing field '{field}' in element"
def test_confidence_values_in_range(self, processor, text_drawing):
result = processor.extract_text(text_drawing)
for el in result["elements"]:
assert 0 <= el["confidence"] <= 100
def test_detects_dimension_number(self, processor, text_drawing):
result = processor.extract_text(text_drawing)
all_texts = [el["text"] for el in result["elements"]]
assert any("3640" in t for t in all_texts), \
f"Expected '3640' in OCR output, got: {all_texts}"
def test_classifies_bare_number_as_number(self, processor, text_drawing):
result = processor.extract_text(text_drawing)
num_texts = [el["text"] for el in result["classified"]["numbers"]]
assert "3640" in num_texts, \
f"'3640' (no unit) should be in numbers, got: {num_texts}"
def test_blank_image_returns_empty(self, processor):
blank = np.ones((400, 600), dtype=np.uint8) * 255
result = processor.extract_text(blank)
assert result["element_count"] == 0
assert result["elements"] == []
def test_full_text_is_string(self, processor, synthetic_drawing):
result = processor.extract_text(synthetic_drawing)
assert isinstance(result["full_text"], str)
class TestClassify:
def test_dimension_pattern(self, processor):
elements = [{"text": "3640mm", "confidence": 90, "x": 0, "y": 0, "width": 50, "height": 20}]
classified = processor._classify(elements)
assert any(e["text"] == "3640mm" for e in classified["dimensions"])
def test_room_name_pattern(self, processor):
elements = [{"text": "Kitchen", "confidence": 90, "x": 0, "y": 0, "width": 80, "height": 20}]
classified = processor._classify(elements)
assert any(e["text"] == "Kitchen" for e in classified["room_names"])
def test_number_pattern(self, processor):
elements = [{"text": "42", "confidence": 90, "x": 0, "y": 0, "width": 30, "height": 20}]
classified = processor._classify(elements)
assert any(e["text"] == "42" for e in classified["numbers"])
def test_dimension_with_unit(self, processor):
elements = [{"text": "3640mm", "confidence": 90, "x": 0, "y": 0, "width": 60, "height": 20}]
classified = processor._classify(elements)
assert any(e["text"] == "3640mm" for e in classified["dimensions"])
def test_bare_number_not_in_dimensions(self, processor):
elements = [{"text": "100", "confidence": 90, "x": 0, "y": 0, "width": 40, "height": 20}]
classified = processor._classify(elements)
assert not any(e["text"] == "100" for e in classified["dimensions"])
def test_label_catchall(self, processor):
elements = [{"text": "note", "confidence": 90, "x": 0, "y": 0, "width": 40, "height": 20}]
classified = processor._classify(elements)
assert any(e["text"] == "note" for e in classified["labels"])
|