hr-ocr-schoolcert/core/ocr.py

from dataclasses import dataclass
from core.parser import *
from core.direction import *
import numpy as np
from typing import List
from core.line_parser import OcrResult, LineParser
from paddleocr import PaddleOCR
from utils.time import timeit


@dataclass
class IdCardOcr:
    ocr: PaddleOCR
    angle_detector: AngleDetector

    def predict(self, image: np.ndarray, image_type: str = '1'):
        img_type = int(image_type)

        img, angle, result = self._rotate_img(image, img_type)
        print(f'---------- detect angle: {angle} 图片角度 ----------')

        if angle != 0:
            _, _, result = self._ocr(img)

        return self._post_process(result, angle, image_type)

    @timeit
    def _rotate_img(self, image, image_type) -> (np.ndarray, int):
        angle, result = self.angle_detector.detect_angle(image, image_type)
        return image, angle, result

    @timeit
    def _ocr(self, image):
        result = self.ocr.ocr(image, cls=True)
        if not result: raise Exception('无法识别')
        confs = [line[1][1] for line in result]

        txts = [line[1][0] for line in result]
        return txts, confs, result

    @timeit
    def _post_process(self, result, angle: int, image_type: str):
        line_parser = LineParser(result)
        line_result = line_parser.parse()
        conf = line_parser.confidence

        # 0:教育部学历证书电子注册备案表
        # 1:教育部学籍在线验证报告
        # 2:中国高等教育证书查询结果(零散查询)

        if int(image_type) in {0, 1, 2}:
            parser = PostParser(line_result)
        else:
            raise Exception('未传入 image_type')
        ocr_res = parser.parse()
        res = {
            "confidence": conf,
            "img_type": str(image_type),
            "orientation": angle,  # 原angle是逆时针，转成顺时针
            **ocr_res
        }

        print(res)
        return res