OCR/blues/idc.py

from flask import Blueprint, views, render_template, request
from utils.util import *
import re
from utils.conf import MAX_CONTENT_LENGTH

idc = Blueprint("idc", __name__, url_prefix="/idc")

__name_ptn = r"姓 *名 *(?P<name>.+) *"
__gender_nation_ptn = r"性 *别 *(?P<gender>男|女) *民 *族 *(?P<nation>.+) *"
__birth_ymd_ptn = r"出 *生 *(?P<year>\d{4}) *年 *(?P<month>\d{2}) *月 *(?P<day>\d{2}) *日 *"
__addr_start_ptn = r"住 *址 *(?P<addr>.+) *"
__idn_ptn = r"公 *民 *身 *份 *号 *码 *(?P<idn>\d{18}) *"
__agent_ptn = r"签 *发 *机 *关 *(?P<agent>.*) *"
__valid_date_ptn = r"有 *效 *期 *限 *(?P<from_year>\d{4})\.(?P<from_month>\d{2})\.(?P<from_day>\d{2})" \
                   r"[^\d]+(?P<to_year>\d{4})\.(?P<to_month>\d{2})\.(?P<to_day>\d{2}) *"


def get_face_info(data: "list[str]"):
    res = {"name": "", "gender": "", "nation": "", "addr": "", "idn": "", "birth": {"year": "", "month": "", "day": ""}}
    for item in data:
        if name := re.match(__name_ptn, item):
            res["name"] = name.group("name")
        elif gender_nation := re.match(__gender_nation_ptn, item):
            res["gender"] = gender_nation.group("gender")
            res["nation"] = gender_nation.group("nation")
        elif birth_ymd := re.match(__birth_ymd_ptn, item):
            res["birth"]["year"] = birth_ymd.group("year")
            res["birth"]["month"] = birth_ymd.group("month")
            res["birth"]["day"] = birth_ymd.group("day")
        elif addr := re.match(__addr_start_ptn, item):
            res["addr"] = addr.group("addr")
        elif idn := re.match(__idn_ptn, item):
            res["idn"] = idn.group("idn")
        else:
            res["addr"] += item
    return res


def get_icon_info(data: "list[str]"):
    res = {"agent": "", "from": {"year": "", "month": "", "day": ""}, "to": {"year": "", "month": "", "day": ""}}
    for item in data:
        if agent := re.match(__agent_ptn, item):
            res["agent"] = agent.group("agent")
        elif valid_date := re.match(__valid_date_ptn, item):
            res["from"]["year"] = valid_date.group("from_year")
            res["from"]["month"] = valid_date.group("from_month")
            res["from"]["day"] = valid_date.group("from_day")
            res["to"]["year"] = valid_date.group("to_year")
            res["to"]["month"] = valid_date.group("to_month")
            res["to"]["day"] = valid_date.group("to_day")
    return res


class IdcView(views.MethodView):
    @staticmethod
    def get():
        return render_template("idc_index.html")

    @staticmethod
    def post():
        pic = request.files.get("picture")
        if pic is None:
            return Response("empty body")
        ext = get_ext_name(pic.filename)
        if not is_image_ext(ext):
            return Response("文件类型错误")
        content = pic.read()
        if len(content) > MAX_CONTENT_LENGTH:
            return Response("文件过大，请重新选择")
        raw_path = f"static/images/{current_time()}_{rand_str()}.{ext}"
        with open(raw_path, "wb") as fp:
            fp.write(content)
            fp.close()

        which = request.form.get("which")
        if which is not None:
            which = which.lower()
        if which not in ["face", "icon"]:
            return Response(f"not recognized arg <which>: '{which}'")

        ocr_res, _ = recognize(content)
        words = [it[1][0] for it in ocr_res]
        if which == "face":
            info = get_face_info(words)
            if json_all(info):
                return Response(data=info)
            return Response("识别失败，请重新选择", info)
        info = get_icon_info(words)
        if json_all(info):
            return Response(data=info)
        return Response("识别失败，请重新选择", info)


class IdcHtmlView(views.MethodView):
    @staticmethod
    def post():
        pic = request.files.get("picture")
        if pic is None:
            return Response("empty body")
        ext = get_ext_name(pic.filename)
        if not is_image_ext(ext):
            return Response("文件类型错误")
        content = pic.read()
        if len(content) > MAX_CONTENT_LENGTH:
            return Response("文件过大，请重新选择")
        cut, rnd = current_time(), rand_str()
        raw_path = f"static/images/{cut}_{rnd}.{ext}"
        rec_path = f"static/images/{cut}_{rnd}_rec.{ext}"
        with open(raw_path, "wb") as fp:
            fp.write(content)
            fp.close()

        which = request.form.get("which")
        if which is not None:
            which = which.lower()
        if which not in ["face", "icon"]:
            return Response(f"not recognized arg <which>: '{which}'")

        ocr_res, img_shape = recognize(content)
        words = [it[1][0] for it in ocr_res]
        draw_img(img_shape, [{"pos": it[0], "word": it[1][0], "rate": it[1][1]} for it in ocr_res], rec_path)
        if which == "face":
            info = get_face_info(words)
        else:
            info = get_icon_info(words)
        info["SUCCESS"] = str(json_all(info)).upper()
        return render_template("k-v_result.html", raw=raw_path, rec=rec_path, data=info)


idc.add_url_rule("/", view_func=IdcView.as_view("idc"))
idc.add_url_rule("/html/", view_func=IdcHtmlView.as_view("idc-html"))