hr-ocr-idcard/core/line_parser.py

import numpy as np
from dataclasses import dataclass

# result 对象
@dataclass
class OcrResult(object):
    box: np.ndarray
    txt: str
    conf: float

    def __hash__(self):
        return hash(repr(self))

    def __repr__(self):
        return f'txt: {self.txt}, box: {self.box.tolist()}, conf: {self.conf}'

    @property
    def lt(self):
        l, t = np.min(self.box, 0)
        return [l, t]

    @property
    def rb(self):
        r, b = np.max(self.box, 0)
        return [r, b]

    @property
    def wh(self):
        l, t = self.lt
        r, b = self.rb
        return [r - l, b - t]

    @property
    def center(self):
        l, t = self.lt
        r, b = self.rb
        return [(r + l)/2, (b + t)/2]

    def one_line(self, b, is_horizontal, eps: float = 20.0) -> bool:
        if is_horizontal:
            return abs(self.lt[1] - b.lt[1]) < eps
        else:
            return abs(self.rb[0] - b.rb[0]) < eps


# 行处理器
class LineParser(object):
    def __init__(self, ocr_raw_result):
        self.ocr_res = []
        for re in ocr_raw_result:
            o = OcrResult(np.array(re[0]), re[1][0], re[1][1])
            self.ocr_res.append(o)
        self.eps = self.avg_height * 0.66

    @property
    def is_horizontal(self):
        res = self.ocr_res
        wh = np.stack([np.abs(np.array(r.lt) - np.array(r.rb)) for r in res])
        return np.sum(wh[:, 0] > wh[:, 1]) > np.sum(wh[:, 0] < wh[:, 1])

    @property
    def avg_height(self):
        idx = self.is_horizontal + 0
        return np.mean(np.array([r.wh[idx] for r in self.ocr_res]))

    # 整体置信度
    @property
    def confidence(self):
        return np.mean([r.conf for r in self.ocr_res])

    # 处理器函数
    def parse(self, eps=40.0):
        # 存返回值
        res = []

        # 需要 处理的 OcrResult 对象  的长度
        length = len(self.ocr_res)

        # 如果字段数 小于等于1 就抛出异常
        if length <= 1:
            raise Exception('无法识别')

        # 遍历数组 并处理他
        for i in range(length):
            # 拿出 OcrResult对象的 第i值 -暂存-
            res_i = self.ocr_res[i]

            # any：-> True
            #       -input: 可迭代对象   |   -output: bool
            #       -如果iterable的任何元素为true，则返回true。如果iterable为空，则返回false。 -与🚪-
            # map: -> [False, False, False, False, True, True, False, False]
            #       -input: (函数, 可迭代对象)     ｜    -output: 可迭代对象
            #       -把 res 喂给lambda --lambda返回True的值-->  输出 新的可迭代对象

            # 这次的 res_i 之前已经在结果集中，就继续下一个
            if any(map(lambda x: res_i in x, res)): continue

            # set() -> {}
            # 初始化一个集合 即-输出-
            res_row = set()

            for j in range(i, length):
                res_j = self.ocr_res[j]
                if res_i.one_line(res_j, self.is_horizontal, self.eps):
                    # LineParser 对象  不可以直接加入字典

                    res_row.add(res_j)
            res.append(res_row)
        idx = self.is_horizontal + 0
        return sorted([list(r) for r in res], key=lambda x: x[0].lt[idx])