OCR/blues/bank.py

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

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

_NumberPtn = re.compile(r"(?P<number>\d{16,19})")


class BankView(views.MethodView):
    @staticmethod
    def get():
        return render_template("bank_index.html")

    @staticmethod
    def post():
        start = time()
        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("文件过大，请重新选择")

        img = read_img(content)
        cropped = crop_img(img)  # 边缘裁剪，对深色背景的效果很好
        images = rot_img_2(cropped)
        recognizes = Engine.rec_multi(images)
        info, sta, idx = {}, False, 0
        for i, ocr_res in enumerate(recognizes):
            rec_str = "".join([it[0] for it in ocr_res])
            if search := _NumberPtn.search(rec_str):
                sta, idx = True, i
                info["number"] = search.group("number")
                break
            else:
                Logger.error(rec_str)
        info["duration"] = time() - start
        if sta:
            raw_path = f"static/images/{current_time()}_{rand_str()}.{ext}"
            save_img(raw_path, images[idx])
            return Response(data=info)
        return Response("识别失败，请重新选择", info)


class BankHtmlView(views.MethodView):
    @staticmethod
    def post():
        start = time()
        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("文件过大，请重新选择")

        img = read_img(content)
        cropped = crop_img(img)
        images = rot_img_2(cropped)
        recognizes = Engine.rec_multi(images)
        info, err_rec, sta, idx = {}, [], False, 0
        msg = "识别失败，请重新选择"
        for i, ocr_res in enumerate(recognizes):
            rec_str = "".join([it[0] for it in ocr_res])
            if search := _NumberPtn.search(rec_str):
                sta, idx = True, i
                info["number"] = search.group("number")
                msg = "识别成功"
                break
            else:
                Logger.error(rec_str)
                err_rec.append(rec_str)

        file_path = f"static/images/{current_time()}_{rand_str()}.{ext}"
        save_img(file_path, images[idx])

        info["SUCCESS"] = str(sta).upper()
        info["MESSAGE"] = msg
        if len(err_rec):
            info["MESSAGE"] += "<br>识别结果：<br>" + "<br>".join(err_rec)
        info["DURATION"] = time() - start  # noqa
        return render_template("k-v_result.html", raw=file_path, data=info)


bank.add_url_rule("/", view_func=BankView.as_view("bank"))
bank.add_url_rule("/html/", view_func=BankHtmlView.as_view("bank-html"))