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)