OCR/utils/util.py

import cv2
import numpy as np
from flask import jsonify
from paddleocr import PaddleOCR
from random import randint, seed
from time import localtime, strftime
from paddleocr.tools.infer.utility import draw_box_txt_fine

__all__ = [
    "Args", "Response", "rand_str", "current_time", "get_ext_name", "is_image_ext", "recognize", "draw_img"
]

__StrBase = "qwertyuioplkjhgfdsazxcvbnm1234567890ZXCVBNMLKJHGFDSAQWERTYUIOP"
__StrBaseLen = len(__StrBase) - 1
__AcceptExtNames = ["jpg", "jpeg", "bmp", "png", "rgb", "tif", "tiff", "gif", "pdf"]
__OcrEngine = PaddleOCR(
    use_gpu=True,
    det_model_dir="models/det/",
    rec_model_dir="models/rec/",
    cls_model_dir="models/cls/",
    use_angle_cls=True,
    use_space_char=True
)


class Args:
    def __init__(self, **kwargs):
        self.__update(
            use_gpu=False, use_xpu=False, use_npu=False, ir_optim=True, use_tensorrt=False,
            min_subgraph_size=15, precision="fp32", gpu_mem=500, image_dir=None, page_num=0,
            det_algorithm="DB", det_model_dir="models/det/", det_limit_side_len=960, det_limit_type="max",
            det_box_type="quad", det_db_thresh=0.3, det_db_box_thresh=0.6, det_db_unclip_ratio=1.5,
            max_batch_size=10, use_dilation=False, det_db_score_mode="fast", det_east_score_thresh=0.8,
            det_east_cover_thresh=0.1, det_east_nms_thresh=0.2, det_sast_score_thresh=0.5,
            det_sast_nms_thresh=0.2, det_pse_thresh=0, det_pse_box_thresh=0.85, det_pse_min_area=16,
            det_pse_scale=1, scales=[8, 16, 32], alpha=1.0, beta=1.0, fourier_degree=5,
            rec_algorithm="SVTR_LCNet", rec_model_dir="models/rec/", rec_image_inverse=True,
            rec_image_shape="3, 48, 320", rec_batch_num=6, max_text_length=25,
            rec_char_dict_path="venv/lib/site-packages/paddleocr/ppocr/utils/ppocr_keys_v1.txt",
            use_space_char=True, vis_font_path="static/simfang.ttf", drop_score=0.5,
            e2e_algorithm="PGNet", e2e_model_dir=None, e2e_limit_side_len=768, e2e_limit_type="max",
            e2e_pgnet_score_thresh=0.5, e2e_char_dict_path="./ppocr/utils/ic15_dict.txt",
            e2e_pgnet_valid_set="totaltext", e2e_pgnet_mode="fast", use_angle_cls=True,
            cls_model_dir="models/cls/", cls_image_shape="3, 48, 192", label_list=["0", "180"],
            cls_batch_num=6, cls_thresh=0.9, enable_mkldnn=True, cpu_threads=10, use_pdserving=False,
            sr_model_dir=None, sr_image_shape="3, 32, 128", sr_batch_num=1,
            draw_img_save_dir="static/rec_res/", save_crop_res=False, crop_res_save_dir="./output",
            use_mp=False, benchmark=False, save_log_path="./log_output/",
            show_log=True, use_onnx=False, output="./output", table_max_len=488, table_algorithm="TableAttn",
            table_model_dir=None, merge_no_span_structure=True, table_char_dict_path=None,
            layout_model_dir=None, layout_dict_path=None, layout_score_threshold=0.5,
            layout_nms_threshold=0.5, kie_algorithm="LayoutXLM", ser_model_dir=None, re_model_dir=None,
            use_visual_backbone=True, ser_dict_path="../train_data/XFUND/class_list_xfun.txt",
            ocr_order_method=None, mode="structure", image_orientation=False, layout=True, table=True,
            ocr=True, recovery=False, use_pdf2docx_api=False, lang="ch", det=True, rec=True, type="ocr",
            ocr_version="PP-OCRv3", structure_version="PP-StructureV2"
        )

        self.__update(**kwargs)

    def __update(self, **kwargs):
        for k, v in kwargs:
            self.__dict__[k] = v

    def __setattr__(self, key: "str", value):
        self.__dict__[key] = value

    def __getattribute__(self, key: "str"):
        assert key in self.__dict__.keys()
        return self.__dict__[key]


def rand_str(size: "int" = 8) -> "str":
    return "".join([__StrBase[randint(0, __StrBaseLen)] for _ in range(size)])


def current_time() -> "str":
    return strftime("%Y-%m-%d_%H-%M-%S", localtime())


def get_ext_name(name: "str") -> "str":
    return name.split(".")[-1].lower()


def is_image_ext(ext: "str") -> bool:
    return ext in __AcceptExtNames


def Response(message: "str" = None, data=None):
    if message is None:
        return jsonify(success=True, message="操作成功", data=data)
    return jsonify(success=False, message=message, data=data)


def recognize(content: "str") -> "tuple[list, tuple]":
    img = cv2.imdecode(np.fromstring(content, np.uint8), 1)  # noqa

    return __OcrEngine.ocr(img)[0], img.shape


def draw_img(shape: "tuple", data: "list[dict]", path: "str", drop: "float" = 0.5):
    img = np.ones(shape, dtype=np.uint8) * 255
    seed(0)

    for one in data:
        if one["rate"] < drop:
            continue
        color = (randint(0, 255), randint(0, 255), randint(0, 255))
        text = draw_box_txt_fine((shape[1], shape[0]), one["pos"], one["word"], font_path="static/simfang.ttf")
        pts = np.array(one["pos"], np.int32).reshape((-1, 1, 2))
        cv2.polylines(text, [pts], True, color, 1)  # noqa
        img = cv2.bitwise_and(img, text)  # noqa

    cv2.imwrite(path, np.array(img))  # noqa