hr-ocr-regbook/core/parser.py

import re
from dataclasses import dataclass
from collections import defaultdict
from typing import List
from core.line_parser import OcrResult
import numpy as np
import cpca
import address_correction.fix_address as fa
@dataclass
class RecItem:
    text: str = ''
    confidence: float = 0.

    def to_dict(self):
        return {"text": self.text, "confidence": np.nan_to_num(self.confidence)}



# 父类
class Parser(object):
    def __init__(self, ocr_results: List[List[OcrResult]], ocr_line: List[List[OcrResult]]):
        self.result = ocr_results
        self.result_line = ocr_line
        self.bool_ranks = bool(ocr_results)
        self.id_index = None
        self.res = defaultdict(RecItem)
        self.keys = ['type', "address", 'address_province', 'address_city', 'address_region', 'address_detail',
                     'name', 'id', 'gender',
                     # 出生地
                     'birthplace', 'birthplace_province', 'birthplace_city', 'birthplace_region',
                     # 籍贯
                     'native_place', 'native_place_province', 'native_place_city', 'native_place_region',
                     'blood_type', 'religion']
        for key in self.keys:
            self.res[key] = RecItem()

        ch = re.compile(u'[\u4e00-\u9fa5+\u0030-\u0039\u0041-\u005a\u0061-\u007a]')
        if self.bool_ranks:
            for item in self.result:
                tail = ['', 1.]
                for k in range(len(item)):
                    item[k].txt = ''.join(re.findall(ch, item[k].txt))
                    tail[0] = tail[0] + item[k].txt
                    tail[1] = tail[1] + item[k].conf
                tail[1] = (tail[1] - 1.) / len(item)
                item.append(tail)
        for item in self.result_line:
            tail = ['', 1.]
            for k in range(len(item)):
                item[k].txt = ''.join(re.findall(ch, item[k].txt))
                tail[0] = tail[0] + item[k].txt
                tail[1] = tail[1] + item[k].conf
            tail[1] = (tail[1] - 1.) / len(item)
            item.append(tail)

    def parse(self):
        return self.res

    def split_addr(self, place: str):

        if place == "birth":
            place = "birthplace"
        elif place == "native":
            place = "native_place"
        elif place == "address":
            place = "address"

        print(self.res[place].text, '=======')
        conf = self.res[place].confidence
        df = cpca.transform([self.res[place].text])
        if df.adcode[0] is None:
            self.res[f"{place}_detail"] = RecItem(self.res[place].text, self.res[place].confidence)
            return
        df = df.replace([None], [''])

        province = df.iloc[0, 0] or ''
        city = df.iloc[0, 1] or ''
        region = df.iloc[0, 2] or ''
        detail = df.iloc[0, 3] or ''
        # 修复地名更新
        if len(detail) == 1: detail = ''
        "===========地址纠错============"
        province = fa.fix_first_level(province)
        # 1. 当province有值、city为空、region为空，detail为空不走地址纠错 即 '上海市'
        # 2. 当province有值、city有值、region有值，不走地址纠错 即 '三级行政单位都识别出来了'
        unfix_add_1 = province != '' and city == '' and region == '' and detail == ''
        unfix_add_2 = province != '' and city != '' and region == '' and detail == ''
        unfix_add_3 = province != '' and city != '' and region != ''

        if (unfix_add_1 or unfix_add_2 or unfix_add_3) is False:
            if place == "address":
                # 1. 第二级单位未识别出来 ---- 切三个字符去纠错
                if city == '':
                    if province == "内蒙古自治区":
                        may_region = detail
                    elif '县' in detail:
                        may_region = detail.split('县')[0] + '县'
                    elif '区' in detail:
                        may_region = detail.split('区')[0] + '区'
                    elif '族' in detail:
                        may_region = detail.split('族')[0] + '族'
                    else:
                        may_region = detail[:3]

                    province, city = fa.fix_second_level(province, may_region)

                    if city in detail:
                        detail = detail[len(city):]
                    province, city, region = fa.fix_third_level(province, city, may_region)
                    detail = detail[len(region):]
                if region == '':
                    province, city, region = fa.fix_third_level(province, city, detail if province == "内蒙古自治区" or '族' in detail else detail[:3])
                    detail = detail[len(region):]
            elif city == '':
                if province == "内蒙古自治区":
                    may_region = detail
                elif '市' in detail:
                    may_region = detail.split('市')[0] + '市'
                elif '族' in detail:
                    may_region = detail.split('族')[0] + '族'
                elif '县' in detail:
                    may_region = detail.split('县')[0] + '县'
                elif '区' in detail:
                    may_region = detail.split('区')[0] + '区'
                else:
                    may_region = detail[:3]

                province, city = fa.fix_second_level(province, may_region)

                if city in detail:
                    detail = detail[len(city):]
                province, city, region = fa.fix_third_level(province, city, may_region)
                detail = detail[len(region):]
            elif region == '':
                province, city, region = fa.fix_third_level(province, city, detail if province == "内蒙古自治区" or '族' in detail else detail[:3])
                detail = ''
            else:
                detail = ''
                province, city, region = fa.fix_third_level(province, city, region)

        "===========地址纠错============"
        print(f'pronvince: {province}, city: {city}, region: {region}, detail: {detail}')
        self.res[f"{place}_province"] = RecItem(province, conf)
        self.res[f"{place}_city"] = RecItem(city, conf)
        if detail and "旗" in detail and "红旗" not in detail:
            temp_region = []
            temp_region.insert(0, detail.split("旗")[0] + "旗")
            self.res[f"{place}_region"] = RecItem(temp_region[0], conf)
            self.res[f"{place}_detail"] = RecItem(detail.split("旗")[-1], conf)


        else:
            self.res[f"{place}_region"] = RecItem(region, conf)
            self.res[f"{place}_detail"] = RecItem(detail, conf)

        self.res[place].text = province + city + region + detail


# 1 户口本首页
class FrontRegBookParser(Parser):
    def type_(self):
        """
        户别
        """
        def get_txt(bool_ranks):
            if bool_ranks:
                gtxt = fa.Correction(self.result[1][0].txt)
                gconf = self.result[1][0].conf
                return gtxt, gconf
            else:
                for i in self.result_line[:-1]:
                    if '户别' in i[-1][0] or i[-1][0][0] == '别':
                        gtxt = i[-1][0].split('别')[-1].split('户主')[0]
                        gconf = i[-1][1]
                        gtxt = re.sub(r'[0-9]+', '', gtxt)
                        return gtxt, gconf

        txt, conf = get_txt(self.bool_ranks)
        if txt == '':
            txt, conf = get_txt(False)

        self.res["type"] = RecItem(fa.Correction(txt), conf)
        return


    def address(self):
        """
        首页住址
        """
        address_txt = ''
        address_conf = 0.
        if self.bool_ranks:
            address = self.result[0][0]
            address_txt = address.txt
            address_conf = address.conf

        else:
            for i in self.result_line[:-1]:
                if '住址' in i[-1][0]:
                    address_txt = i[-1][0].split('住址')[-1]
                    address_conf = i[-1][1]

        self.res["address"] = RecItem(fa.Correction(address_txt), address_conf)
        self.split_addr('address')

    # 存入
    def parse(self):
        self.type_()
        self.address()
        return {key: self.res[key].to_dict() for key in self.keys}


# 0 常驻人口页
class PeopleRegBookParser(Parser):

    def full_name(self):
        """
        姓名
        属 result[1]
        位 0
        """
        name = self.result[1][0]
        name_val = fa.Correction(name.txt)
        conf = name.conf
        if len(name_val) < 5:
            self.res["name"] = RecItem(name_val, conf)
        else:
            point_unicode = ["\u2E31", "\u2218", "\u2219", "\u22C5", "\u25E6", "\u2981",
                             "\u00B7", "\u0387", "\u05BC", "\u16EB", "\u2022", "\u2027",
                             "\u2E30", "\uFF0E", "\u30FB", "\uFF65", "\u10101"]
            for n in range(len(point_unicode)):
                point = re.findall(point_unicode[n], name_val)
                if len(point) != 0:
                    name_list = name_val.split(point[0])
                    self.res['name'] = RecItem(name_list[0] + '\u00B7' + name_list[1], conf)
                    return
                else:
                    self.res["name"] = RecItem(name_val, conf)
                    return

    def card_no(self):
        """
        身份证号码
        属 0/1都可
        """
        for i in range(len(self.result[1][:-1])):
            res = self.result[1][i]
            txt = res.txt
            conf = res.conf
            if "X" in txt or "x" in txt:
                code_val = re.findall("\d*[X|x]", txt)
            else:
                code_val = re.findall("\d{10,18}", txt)
            if len(code_val) > 0:
                id_num = fa.Correction(str(code_val[0]))
                self.id_index = i
                if len(id_num) == 18 or len(id_num) > 10:
                    self.res["id"].text = id_num
                    self.res["id"].confidence = conf
                    if len(id_num) == 18:
                        self.res["gender"].text = ("男" if int(id_num[16]) % 2 else "女")
                    self.res["gender"].confidence = conf
                return
        raise Exception('身份证号识别出错')

    def blood_type(self):
        """
        血型
        属 result[1]
        """

        blood_val = ''
        conf = 0.
        for res in self.result[1][:-1]:
            txt = res.txt
            conf = res.conf
            if "血型" in txt:
                blood_val = txt.split("血型")[-1]
            self.res["blood_type"] = RecItem(blood_val, conf)

    def religion(self):
        """
        宗教信仰
        """

        religion_val = ''
        for res in self.result[1][:-1]:
            txt = res.txt
            conf = res.conf
            if "宗教信仰是" in txt:
                religion_val = txt.split("宗教信仰是")[-1]
            self.res["religion"] = RecItem(fa.Correction(religion_val), conf)

    def birthplace(self):
        """
        出生地
        """
        birth_place = self.result[1][1]
        birth_place_txt = birth_place.txt
        birth_place_conf = birth_place.conf

        if birth_place_txt:
            self.res["birthplace"] = RecItem(fa.Correction(birth_place_txt), birth_place_conf)
            self.split_addr("birth")

    def native_place(self):
        """
        籍贯
        """
        native_place = self.result[1][2]
        native_place_txt = native_place.txt
        native_place_conf = native_place.conf
        self.res["native_place"] = RecItem(fa.Correction(native_place_txt), native_place_conf)
        self.split_addr("native")



    def parse(self):
        self.full_name()
        self.card_no()
        self.blood_type()
        self.religion()
        self.birthplace()
        self.native_place()


        return {k: self.res[k].to_dict() for k in self.keys}