reader.py

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