You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
4 weeks ago
|
import os
|
||
|
|
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from PIL import Image, ImageDraw
|
||
|
|
|
||
|
|
from surya.detection import DetectionPredictor, InlineDetectionPredictor
|
||
|
|
from surya.ocr_error import OCRErrorPredictor
|
||
|
|
from surya.layout import LayoutPredictor
|
||
|
|
from surya.recognition import RecognitionPredictor
|
||
|
|
from surya.table_rec import TableRecPredictor
|
||
|
|
from surya.texify import TexifyPredictor
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture(scope="session")
|
||
|
|
def ocr_error_predictor() -> OCRErrorPredictor:
|
||
|
|
ocr_error_predictor = OCRErrorPredictor()
|
||
|
|
yield ocr_error_predictor
|
||
|
|
del ocr_error_predictor
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture(scope="session")
|
||
|
|
def layout_predictor() -> LayoutPredictor:
|
||
|
|
layout_predictor = LayoutPredictor()
|
||
|
|
yield layout_predictor
|
||
|
|
del layout_predictor
|
||
|
|
|
||
|
|
@pytest.fixture(scope="session")
|
||
|
|
def detection_predictor() -> DetectionPredictor:
|
||
|
|
detection_predictor = DetectionPredictor()
|
||
|
|
yield detection_predictor
|
||
|
|
del detection_predictor
|
||
|
|
|
||
|
|
@pytest.fixture(scope="session")
|
||
|
|
def recognition_predictor() -> RecognitionPredictor:
|
||
|
|
recognition_predictor = RecognitionPredictor()
|
||
|
|
yield recognition_predictor
|
||
|
|
del recognition_predictor
|
||
|
|
|
||
|
|
@pytest.fixture(scope="session")
|
||
|
|
def table_rec_predictor() -> TableRecPredictor:
|
||
|
|
table_rec_predictor = TableRecPredictor()
|
||
|
|
yield table_rec_predictor
|
||
|
|
del table_rec_predictor
|
||
|
|
|
||
|
|
@pytest.fixture(scope="session")
|
||
|
|
def texify_predictor() -> TexifyPredictor:
|
||
|
|
texify_predictor = TexifyPredictor()
|
||
|
|
yield texify_predictor
|
||
|
|
del texify_predictor
|
||
|
|
|
||
|
|
@pytest.fixture(scope="session")
|
||
|
|
def inline_detection_predictor() -> InlineDetectionPredictor:
|
||
|
|
inline_detection_predictor = InlineDetectionPredictor()
|
||
|
|
yield inline_detection_predictor
|
||
|
|
del inline_detection_predictor
|
||
|
|
|
||
|
|
@pytest.fixture()
|
||
|
|
def test_image():
|
||
|
|
image = Image.new("RGB", (1024, 1024), "white")
|
||
|
|
draw = ImageDraw.Draw(image)
|
||
|
|
draw.text((10, 10), "Hello World", fill="black", font_size=72)
|
||
|
|
draw.text((10, 200), "This is a sentence of text.\nNow it is a paragraph.\nA three-line one.", fill="black",
|
||
|
|
font_size=24)
|
||
|
|
return image
|
||
|
|
|
||
|
|
@pytest.fixture()
|
||
|
|
def test_image_tall():
|
||
|
|
image = Image.new("RGB", (4096, 4096), "white")
|
||
|
|
draw = ImageDraw.Draw(image)
|
||
|
|
draw.text((10, 10), "Hello World", fill="black", font_size=72)
|
||
|
|
draw.text((4000, 4000), "This is a sentence of text.\n\nNow it is a paragraph.\n\nA three-line one.", fill="black", font_size=24)
|
||
|
|
return image
|
||
|
|
|