OCR/utils/util.py

import cv2
import numpy as np
from typing import Union
from flask import jsonify
from random import randint
from hmOCR import HuiMvOCR, Args
from time import localtime, strftime

__all__ = [
    "Response", "rand_str", "current_time", "get_ext_name", "is_image_ext",
    "json_all", "str_include", "read_img", "rot_img", "save_img", "Engine"
]

__StrBase = "qwertyuioplkjhgfdsazxcvbnm1234567890ZXCVBNMLKJHGFDSAQWERTYUIOP"
__StrBaseLen = len(__StrBase) - 1
__AcceptExtNames = ["jpg", "jpeg", "bmp", "png", "rgb", "tif", "tiff", "gif"]
Engine = HuiMvOCR(Args())


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 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 json_all(data: "Union[list, dict]") -> "bool":
    if isinstance(data, list):
        for item in data:
            if isinstance(item, str) and not item:
                return False
            elif isinstance(item, (list, dict)) and not json_all(item):
                return False
        return True
    elif isinstance(data, dict):
        for value in data.values():
            if isinstance(value, str) and not value:
                return False
            elif isinstance(value, (list, dict)) and not json_all(value):
                return False
        return True
    raise TypeError(f"except node type are: [list, dict], but got a {type(data)} instead.")


def str_include(str_long: "str", str_short: "str") -> "bool":
    for it in str_short:
        if it not in str_long:
            return False
    return True


def read_img(content: "str") -> "np.ndarray":
    return cv2.imdecode(np.frombuffer(content, np.uint8), 1)  # noqa


def rot_img(img: "np.ndarray") -> "list[np.ndarray]":
    return [img, np.rot90(img), np.rot90(img, 2), np.rot90(img, 3)]


def save_img(filename: "str", content: "Union[bytes, np.ndarray]"):
    if isinstance(content, np.ndarray):
        return cv2.imwrite(filename, content)  # noqa
    with open(filename, "wb") as fp:
        fp.write(content)
        fp.close()