OCR/utils/util.py

import cv2
import numpy as np
from PIL import Image
from io import BytesIO
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", "compress_img", "crop_img", "rot_img_2", "rot_img_4",
    "save_img", "Engine"
]

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


def __isWhiteBack(img: "np.ndarray") -> "bool":
    row = np.concatenate([img[:__BorderWidth], img[-__BorderWidth:]], axis=0).reshape((__BorderWidth * 2, -1))
    col = np.concatenate([img[:, :__BorderWidth], img[:, -__BorderWidth:]], axis=1).reshape((__BorderWidth * 2, -1))
    full = np.concatenate([row, col], axis=1)
    return np.average(full) > 127


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 compress_img(img: "np.ndarray", tar_size: "int" = __MaxImageSize) -> "np.ndarray":
    cur_size = img.shape[0] * img.shape[1]
    if cur_size <= tar_size:
        return img

    image = Image.fromarray(img)
    output = BytesIO()
    image.save(output, format="JPEG", quality=int(tar_size / cur_size * 100))
    output.seek(0)
    compressed_image = Image.open(output)

    return np.array(compressed_image)


def crop_img(image: "np.ndarray") -> "np.ndarray":
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # noqa 将图像转换为灰度图像
    thresh_type = cv2.THRESH_BINARY  # noqa
    if __isWhiteBack(gray):
        thresh_type = cv2.THRESH_BINARY_INV  # noqa
    _, threshold = cv2.threshold(gray, 155, 255, thresh_type)  # noqa 换为二值图像 => save: [150,255]
    contours, _ = cv2.findContours(threshold, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # noqa 查找轮廓
    if not contours:
        return image
    max_contour = max(contours, key=cv2.contourArea)  # noqa 找到最大的轮廓
    max_area = cv2.contourArea(max_contour)  # noqa
    if max_area * 3 < image.shape[0] * image.shape[1]:
        return image
    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_2(img: "np.ndarray") -> "list[np.ndarray]":
    if img.shape[0] > img.shape[1]:  # 0: height, 1: width
        return [np.rot90(img), np.rot90(img, 3)]
    return [img, np.rot90(img, 2)]


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