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")

__CN = "中国CHINA"
__face_ptn = r"^姓名(?P<name>.+)性别(?P<gender>男|女)民族(?P<nation>.+)" \
             r"出生(?P<year>\d{4})年(?P<month>\d\d)月(?P<day>\d\d)日" \
             r"住址(?P<addr>.+)公民身份号码(?P<idn>\d{17}\d|x|X)$"
__icon_ptn = r"^中华人民共和国居民身份证签发机关(?P<agent>.+)" \
             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})$"


# 需要图片在PC上看着是：横长竖宽
def get_face_info(data: "list[str]") -> "tuple[dict, str, bool]":
    res = {"name": "", "gender": "", "nation": "", "birth": {"year": "", "month": "", "day": ""}, "addr": "", "idn": ""}

    if len(data) < 5:  # 最少 5 个识别结果
        return res, "请使用正确的身份证人像面照片", False
    deal = [item.replace(" ", "") for item in data if not str_include(__CN, item)]
    if not deal[0].startswith("姓名"):  # 非正，逆序后尝试
        deal.reverse()
    if not deal[0].startswith("姓名"):
        return res, "请确保照片为：横长竖宽，正面朝上", False

    str_all = "".join(deal)
    print(str_all)
    if match := re.match(__face_ptn, str_all):
        res["name"] = match.group("name")
        res["gender"] = match.group("gender")
        res["nation"] = match.group("nation")
        res["birth"] = {
            "year": match.group("year"),
            "month": match.group("month"),
            "day": match.group("day")
        }
        res["addr"] = match.group("addr")
        res["idn"] = match.group("idn")
        return res, "", True

    return res, "识别失败，请重新选择", False


def get_icon_info(data: "list[str]"):
    res = {"agent": "", "from": {"year": "", "month": "", "day": ""}, "to": {"year": "", "month": "", "day": ""}}

    if len(data) < 4:  # 最少 4 个识别结果
        return res, "请使用正确的身份证国徽面照片", False
    deal = [item.replace(" ", "") for item in data if not str_include(__CN, item)]
    if not deal[0].startswith("中华"):  # 非正，逆序后尝试
        deal.reverse()
    if not deal[0].startswith("中华"):
        return res, "请确保照片为：横长竖宽，正面朝上", False

    str_all = "".join(deal)
    print(str_all)
    if match := re.match(__icon_ptn, str_all):
        res["agent"] = match.group("agent")
        res["from"] = {
            "year": match.group("from_year"),
            "month": match.group("from_month"),
            "day": match.group("from_day"),
        }
        res["to"] = {
            "year": match.group("to_year"),
            "month": match.group("to_month"),
            "day": match.group("to_day"),
        }
        return res, "", True
    return res, "识别失败，请重新选择", False


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, msg, sta = get_face_info(words)
            if sta:
                return Response(data=info)
            return Response(msg, info)
        info, msg, sta = get_icon_info(words)
        if sta:
            return Response(data=info)
        return Response(msg, 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, msg, sta = get_face_info(words)
        else:
            info, msg, sta = get_icon_info(words)
        info["SUCCESS"] = str(sta).upper()
        info["MESSAGE"] = msg
        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"))