hr-ocr-regbook/core/direction.py

import re

import cv2
import numpy as np
from dataclasses import dataclass
from enum import Enum
from typing import Tuple, List

import cv2
from paddleocr import PaddleOCR

from core.line_parser import LineParser
import matplotlib.pyplot as plt


# 枚举
class Direction(Enum):
    TOP = 0
    RIGHT = 1
    BOTTOM = 2
    LEFT = 3


# 父类
class OcrAnchor(object):
    # anchor的名字, 如身份证号、承办人等
    def __init__(self, name: str, d: List[Direction]):
        self.name = name
        self.direction = d

        # 定义枚举字典
        def t_func(anchor, c, is_horizontal):
            if is_horizontal:
                return 0 if anchor[1] < c[1] else 2
            else:
                return 1 if anchor[0] > c[0] else 3

        def l_func(anchor, c, is_horizontal):
            if is_horizontal:
                return 0 if anchor[0] < c[0] else 2
            else:
                return 1 if anchor[1] < c[1] else 3

        def b_func(anchor, c, is_horizontal):
            if is_horizontal:
                return 0 if anchor[1] > c[1] else 2
            else:
                return 1 if anchor[0] < c[0] else 3

        def r_func(anchor, c, is_horizontal):
            if is_horizontal:
                return 0 if anchor[0] > c[0] else 2
            else:
                return 1 if anchor[1] > c[1] else 3

        self.direction_funcs = {
            Direction.TOP: t_func,
            Direction.LEFT: l_func,
            Direction.BOTTOM: b_func,
            Direction.RIGHT: r_func
        }

    # pic中心点
    def get_pic_center(self, res) -> Tuple[float, float]:
        boxs = []
        for row in res:
            for r in row:
                boxs.extend(r.box)
        boxs = np.stack(boxs)
        l, t = np.min(boxs, 0)
        r, b = np.max(boxs, 0)
        return (l + r) / 2, (t + b) / 2

    # 是否有锚点
    def is_anchor(self, txt, box):
        pass

    # 找锚点
    def find_anchor(self, res):
        for row in res:
            for r in row:
                if self.is_anchor(r.txt, r.box):
                    # l, t = np.min(r.box, 0)
                    # r, b = np.max(r.box, 0)
                    # return True, (l + r) / 2, (t + b) / 2
                    return True, r.center[0], r.center[1]
        return False, 0., 0.

    # get angle
    def locate_anchor(self, res, is_horizontal):
        found, a_cx, a_cy = self.find_anchor(res)
        cx, cy = self.get_pic_center(res)

        if found is False: raise Exception(f'识别不到anchor{self.name}')

        pre = None
        for d in self.direction:
            angle_func = self.direction_funcs.get(d, None)
            angle = angle_func((a_cx, a_cy), (cx, cy), is_horizontal)
            if pre is None:
                pre = angle
            else:
                if pre != angle:
                    raise Exception('angle is not compatible')
        return pre


# 子类1 户口本首页1
class FrontAnchor(OcrAnchor):
    def __init__(self, name: str, d: List[Direction]):
        super(FrontAnchor, self).__init__(name, d)

    def is_anchor(self, txt, box):
        txts = re.findall('承办人', txt)
        if len(txts) > 0:
            return True
        return False

    def locate_anchor(self, res, is_horizontal):
        return super(FrontAnchor, self).locate_anchor(res, is_horizontal)


# 子类2 常驻人口页0
class PeopleAnchor(OcrAnchor):
    def __init__(self, name: str, d: List[Direction]):
        super(PeopleAnchor, self).__init__(name, d)

    def is_anchor(self, txt, box):
        txts = re.findall('常住', txt) or re.findall('登记卡', txt)
        if len(txts) > 0:
            return True
        return False

    def locate_anchor(self, res, is_horizontal):
        return super(PeopleAnchor, self).locate_anchor(res, is_horizontal)


# 调用以上 🔧工具
# <- ocr_生数据
# == ocr_熟数据(行处理后)
# -> 角度0/1/2/3
def detect_angle(result, ocr_anchor: OcrAnchor):
    lp = LineParser(result)
    res = lp.parse()
    print('------ angle ocr -------')
    print(res)
    print('------ angle ocr -------')
    is_horizontal = lp.is_horizontal
    return ocr_anchor.locate_anchor(res, is_horizontal)


@dataclass
class AngleDetector(object):
    """
    角度检测器
    """
    ocr: PaddleOCR

    # 角度检测器
    # <- img(cv2格式)  img_type
    # == result <- img(cv2)
    # -> angle       result(ocr生)
    def detect_angle(self, img, image_type):
        image_type = int(image_type)
        ocr_anchor = PeopleAnchor('常住', [Direction.TOP]) if image_type == 0 else FrontAnchor('承办人', [Direction.BOTTOM,
                                                                                                     Direction.LEFT])

        result = self.ocr.ocr(img, cls=True)

        # image = img.copy()
        # for box in result:
        #     box = np.reshape(np.array(box[0]), [-1, 1, 2]).astype(np.int64)
        #     image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
        # cv2.imwrite("./test.jpg", image)

        try:
            angle = detect_angle(result, ocr_anchor)
            return angle, result


        except Exception as e:
            print(e)
            # 如果第一次识别不到，旋转90度再识别
            img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
            result = self.ocr.ocr(img, cls=True)
            angle = detect_angle(result, ocr_anchor)
            # 旋转90度之后要重新计算角度
            return (angle - 1 + 4) % 4, result