hr-ocr-idcard/core/ocr.py

from dataclasses import dataclass
from core.parser import *
from core.direction import *
import numpy as np
from paddleocr import PaddleOCR


@dataclass
class IdCardOcr:
    ocr: PaddleOCR

    def predict(self, image: np.ndarray, image_type: str = '0'):
        image, angle = self._pre_process(image)
        txts, confs = self._ocr(image)

        if int(image_type) == 0:
            parser = FrontParser(txts, confs)
        elif int(image_type) == 1:
            parser = BackParser(txts, confs)
        else:
            raise Exception('无法识别')
        return self._post_process(angle, parser, image_type)

    def _pre_process(self, image) -> (np.ndarray, int):
        angle = detect_angle(image)
        print(angle)  # 逆时针
        if angle == 180:
            image = cv2.rotate(image, cv2.ROTATE_180)
        if angle == 90:
            image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
        if angle == 270:
            image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
        return image, angle

    def _ocr(self, image):
        # 获取模型检测结果
        result = self.ocr.ocr(image, cls=True)
        print("------------------")
        print(result)
        if not result:
            raise Exception('无法识别')
        confs = [line[1][1] for line in result]

        # 将检测到的文字放到一个列表中
        txts = [line[1][0] for line in result]
        # print("......................................")
        # print(txts)
        # print("......................................")
        return txts, confs

    def _post_process(self, angle: int, parser: Parser, image_type: str):
        content = parser.parse()
        conf = parser.confidence

        res = {
            "confidence": conf,
            "card_type": image_type,
            "orientation": (4 - angle // 90) % 4,  # 原angle是逆时针，转成顺时针
            "name": content["Name"].to_dict(),
            "id": content["IDNumber"].to_dict(),
            "ethnicity": content["Nationality"].to_dict(),
            "gender": content["Gender"].to_dict(),
            "birthday": content["Birth"].to_dict(),
            "address_province": content["address_province"].to_dict(),
            "address_city": content["address_city"].to_dict(),
            "address_region": content["address_region"].to_dict(),
            "address_detail": content["address_detail"].to_dict(),
            "expire_date": content["expire_date"].to_dict()
        }
        print(res)
        return res