from typing import List, Optional, Tuple, Union import easyocr import numpy as np from PIL import Image import asyncio from concurrent.futures import ThreadPoolExecutor from .io import BBOX class ReaderForEasyOCR: def __init__(self, languages: Optional[List[str]] = None, use_gpu: Optional[bool] = None, **kwargs): self.reader = easyocr.Reader(languages, gpu=use_gpu, **kwargs) self.executor = ThreadPoolExecutor() async def __call__(self, inputs: Union[np.ndarray, Image.Image]) -> List[Tuple[BBOX, str]]: if isinstance(inputs, Image.Image): inputs = np.array(inputs) loop = asyncio.get_running_loop() outputs = await loop.run_in_executor(self.executor, self.reader.readtext, inputs) outputs = [ ((int(x1), int(y1), int(x2), int(y2)), text) for ((x1, y1), _, (x2, y2), _), text, _ in outputs ] outputs: List[Tuple[BBOX, str]] = sorted(outputs, key=lambda x: (x[0][1] // 30, x[0][0] // 30)) outputs = list(filter(lambda x: x[1].strip(), outputs)) return outputs