45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
from __future__ import annotations
|
|
|
|
from io import BytesIO
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from PIL import Image
|
|
|
|
from app.local_ocr import LocalOCRError, _sort_reading_order, recognize_image_text
|
|
|
|
|
|
def box(left: int, top: int, right: int, bottom: int):
|
|
return [[left, top], [right, top], [right, bottom], [left, bottom]]
|
|
|
|
|
|
class LocalOcrTests(unittest.TestCase):
|
|
def test_text_detections_are_ordered_by_rows_and_language_direction(self) -> None:
|
|
detections = [
|
|
(box(10, 100, 100, 120), "Place: Amman", 0.9),
|
|
(box(400, 100, 520, 120), "المكان: عمان", 0.8),
|
|
(box(10, 10, 100, 30), "COMMUNITY", 0.9),
|
|
(box(400, 10, 520, 30), "لقاء القراءة", 0.8),
|
|
]
|
|
|
|
ordered = _sort_reading_order(detections)
|
|
|
|
self.assertEqual(
|
|
[text for _box, text, _confidence in ordered],
|
|
["لقاء القراءة", "COMMUNITY", "المكان: عمان", "Place: Amman"],
|
|
)
|
|
|
|
def test_image_pixel_limit_is_checked_before_loading_ocr_model(self) -> None:
|
|
output = BytesIO()
|
|
Image.new("RGB", (1500, 1400), "white").save(output, format="JPEG")
|
|
|
|
with patch("app.local_ocr._get_reader") as get_reader:
|
|
with self.assertRaises(LocalOCRError):
|
|
recognize_image_text(output.getvalue(), label="large test")
|
|
|
|
get_reader.assert_not_called()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|