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",
    "str_include", "read_img", "crop_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 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 crop_img(image: "np.ndarray") -> "np.ndarray":
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # noqa 将图像转换为灰度图像
    _, threshold = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)  # noqa 换为二值图像 => save: [150,255]
    contours, _ = cv2.findContours(threshold, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # noqa 查找轮廓
    max_contour = max(contours, key=cv2.contourArea)  # noqa 找到最大的轮廓
    rect = cv2.minAreaRect(max_contour)  # noqa 计算最小外接矩形
    box = cv2.boxPoints(rect)  # noqa 获取矩形的四个角点
    box = np.intp(box)
    # 裁剪图像
    return image[min(box[:, 1]):max(box[:, 1]), min(box[:, 0]):max(box[:, 0])]


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