hr-ocr-bankcard/core/direction.py

from core.anchor import *
from dataclasses import dataclass
from paddleocr import PaddleOCR


def detect_angle(result, ocr_anchor: OcrAnchor):
    lp = LineParser(result)
    res = lp.parse()
    is_horizontal = lp.is_horizontal
    return ocr_anchor.locate_anchor(res, is_horizontal)


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

    def detect_angle(self, img):
        ocr_anchor = BankCardAnchor('银行卡号', [Direction.BOTTOM])
        result = self.ocr.ocr(img, cls=True)

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

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

    def origin_detect(self, img):
        # 这边一般是在自己的检测模型result=[]时，再使用官方的模型做个检测，如果这个也没有结果，那就真的检测不出来
        result = self.ocr.ocr(img)
        return result

    def det_oneline(self, result):
        # 这边已经是转正之后的图片，不需要考虑是否水平，只要检测是否一行
        lp = LineParser(result)
        return lp.detection_parse()