hr-ocr-bankcard/core/anchor.py

import re
from dataclasses import dataclass
from enum import Enum
from typing import Tuple, List

import cv2
import numpy as np
from paddleocr import PaddleOCR

from core.line_parser import LineParser


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
        # anchor位置
        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.BOTTOM: b_func,
            Direction.LEFT: l_func,
            Direction.RIGHT: r_func,
        }

    # 获取中心区域坐标 -> (x, y)
    def get_rec_area(self, res) -> Tuple[float, float]:
        """获得整张身份证的识别区域, 返回识别区域的中心点"""
        boxes = []
        for row in res:
            for r in row:
                boxes.extend(r.box)
        boxes = np.stack(boxes)
        l, t = np.min(boxes, 0)
        r, b = np.max(boxes, 0)
        # 识别区域的box
        # big_box = [[l, t], [r, t], [r, b], [l, b]]
        # w, h = (r - l, b - t)
        return (l + r) / 2, (t + b) / 2

    # 判断是否是 锚点
    def is_anchor(self, txt, box, conf) -> bool:
        pass

    # 找 锚点 -> 锚点坐标
    def find_anchor(self, res) -> Tuple[bool, float, float]:
        """
        寻找锚点 中心点坐标
        """
        for row in res:
            for r in row:
                txt = r.txt.replace('-', '').replace(' ', '')
                box = r.box
                conf = r.conf
                flag = self.is_anchor(txt, box, conf)
                if flag:
                    l, t = np.min(box, 0)
                    r, b = np.max(box, 0)
                    return True, (l + r) / 2, (t + b) / 2
        #         if flag and (len(re.findall('\d{10,20}', txt)) > 0 and conf > 0.95):
        #             l, t = np.min(box, 0)
        #             r, b = np.max(box, 0)
        #             return True, (l + r) / 2, (t + b) / 2
        #         elif flag:
        #             l, t = np.min(box, 0)
        #             r, b = np.max(box, 0)
        # if l:
        #     return True, (l + r) / 2, (t + b) / 2
        # else:
        return False, 0., 0.

    # 定位 锚点 -> 角度
    # -> 锚点(x, y)  pic(x, y) is_horizontal
    def locate_anchor(self, res, is_horizontal) -> int:
        found, id_cx, id_cy = self.find_anchor(res)

        # 如果识别不到身份证号
        if not found: raise Exception(f'识别不到anchor{self.name}')
        cx, cy = self.get_rec_area(res)
        pre = None
        for d in self.direction:
            f = self.direction_funcs.get(d, None)
            angle = f((id_cx, id_cy), (cx, cy), is_horizontal)
            if pre is None:
                pre = angle
            else:
                if angle != pre:
                    raise Exception('angle is not compatiable')
        return pre


class BankCardAnchor(OcrAnchor):
    def __init__(self, name: str, d: List[Direction]):
        super(BankCardAnchor, self).__init__(name, d)

    def is_anchor(self, txt, box, conf) -> bool:
        # # 这边我动了手脚，可能需要改一下长度，到时候测试再看
        # txts = re.findall('\d{5,20}', txt)
        # # print(txts)
        # if conf > 0.95 and len(txts) > 0:
        #     # print("这是我识别出来的卡号：", txts)
        #     return True

        # 这里逻辑有点长，理想情况下，置信度比较高的txt会在卡号附近，一般在卡号下方
        if len(re.findall('\d{16,20}', txt)) > 0 and conf > 0.95:  # 完美找到卡号
            return True
        elif len(re.findall('\d{10,16}', txt)) > 0 and conf > 0.95:  # 卡号只找到了一半多点
            return True
        elif len(re.findall('\d{6,10}', txt)) > 0 and conf > 0.95:  # 卡号 只找到了一点
            return True
        elif len(re.findall('\d{4,6}', txt)) > 0 and conf > 0.95:  # 卡号只找到了一丢丢
            return True
        elif conf > 0.95:  # 可能卡号就是找到了一个数字，但是置信度很高，
            return True
        # elif conf >= 0.9:
        #     return True

        return False

    def locate_anchor(self, res, is_horizontal) -> int:
        return super(BankCardAnchor, self).locate_anchor(res, is_horizontal)