hr_ocr_general/core/ocr.py

from dataclasses import dataclass
from typing import Any, List
import numpy as np
from paddleocr import PaddleOCR

@dataclass
class OcrRes:
    confidence: float
    text: str
    txt_region: List[List[int]]


@dataclass
class Ocr:
    ocr: PaddleOCR

    def predict(self, image: List[np.ndarray]):
        if ocr_result := self.ocr.ocr(image):
            res = [OcrRes(ocrr[0], ocrr[1][0], ocrr[1][1]) for ocrr in ocr_result]
        else:
            raise Exception('图片中未识别出文字')
        return [res]