import re from dataclasses import dataclass from enum import Enum from typing import Tuple import cv2 import numpy as np from paddleocr import PaddleOCR from core.line_parser import LineParser class Directoin(Enum): TOP = 0 RIGHT = 1 BOTTOM = 2 LEFT = 3 # 父类 class OcrAnchor(object): # 输入识别anchor的名字, 如身份证号 def __init__(self, name: str, d: Directoin): self.name = name # anchor位置 self.direction = d def t_func(anchor, c, is_horizontal): if is_horizontal: return 2 if anchor[1] > c[1] else 0 else: return 3 if anchor[0] < c[0] else 1 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 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 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 = { Directoin.TOP: t_func, Directoin.BOTTOM: b_func, Directoin.LEFT: l_func, Directoin.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) -> 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 if self.is_anchor(txt, box): l, t = np.min(box, 0) r, b = np.max(box, 0) return True, (l + r) / 2, (t + b) / 2 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) # print(f'id_cx: {id_cx}, id_cy: {id_cy}') # print(f'cx: {cx}, cy: {cy}') # 用k->get->func ==> f() f = self.direction_funcs.get(self.direction, None) return f((id_cx, id_cy), (cx, cy), is_horizontal) # if is_horizontal: # # 如果是水平的,身份证号的位置在相对识别区域的下方,方向则为0度,否则是180度 # return 0 if id_cy > cy else 2 # else: # # 如果是竖直的,身份证号的相对位置如果在左边,方向为90度,否则270度 # return 1 if id_cx < cx else 3 # 子类1 人像面 class FrontSideAnchor(OcrAnchor): def __init__(self, name: str, d: Directoin): super(FrontSideAnchor, self).__init__(name, d) def is_anchor(self, txt, box) -> bool: txts = re.findall('\d{10,18}', txt) if len(txts) > 0: return True return False def locate_anchor(self, res, is_horizontal) -> int: return super(FrontSideAnchor, self).locate_anchor(res, is_horizontal) # 子类2 国徽面 class BackSideAnchor(OcrAnchor): def __init__(self, name: str, d: Directoin): super(BackSideAnchor, self).__init__(name, d) def is_anchor(self, txt, box) -> bool: txt = txt.replace('.', '') txts = re.findall('有效期', txt) if len(txts) > 0: return True return False def locate_anchor(self, res, is_horizontal) -> int: return super(BackSideAnchor, 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) # 初始化anchor对象 ocr_anchor = BackSideAnchor('有效期', Directoin.BOTTOM) if image_type != 0 else FrontSideAnchor('身份证号', Directoin.BOTTOM) result = self.ocr.ocr(img, cls=True) 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