48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any, Protocol
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EngineLine:
|
|
text: str
|
|
confidence: float
|
|
bbox: tuple[int, int, int, int]
|
|
|
|
|
|
class OcrEngine(Protocol):
|
|
def recognize(self, image: object) -> list[EngineLine]: ...
|
|
|
|
|
|
class PaddleOcrEngine:
|
|
def __init__(self, pipeline: Any | None = None) -> None:
|
|
self._injected = pipeline is not None
|
|
if pipeline is None:
|
|
from paddleocr import PaddleOCR
|
|
|
|
pipeline = PaddleOCR(
|
|
text_detection_model_name="PP-OCRv5_mobile_det",
|
|
text_recognition_model_name="latin_PP-OCRv5_mobile_rec",
|
|
use_doc_orientation_classify=False,
|
|
use_doc_unwarping=False,
|
|
use_textline_orientation=False,
|
|
device="cpu",
|
|
)
|
|
self.pipeline = pipeline
|
|
|
|
def recognize(self, image: object) -> list[EngineLine]:
|
|
if not self._injected:
|
|
import numpy
|
|
|
|
image = numpy.asarray(image)
|
|
lines: list[EngineLine] = []
|
|
for prediction in self.pipeline.predict(image):
|
|
payload = prediction.json() if callable(prediction.json) else prediction.json
|
|
result = payload.get("res", payload)
|
|
texts = result.get("rec_texts", [])
|
|
scores = result.get("rec_scores", [])
|
|
boxes = result.get("rec_boxes", [])
|
|
if hasattr(boxes, "tolist"):
|
|
boxes = boxes.tolist()
|
|
for text, score, box in zip(texts, scores, boxes, strict=True):
|
|
lines.append(EngineLine(str(text), float(score), tuple(int(value) for value in box)))
|
|
return lines
|