All/Python/PaddleOCR/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", "resize_img", "preprocess_img",
    "rot_img_2", "rot_img_4", "save_img", "Engine"
]

__StrBase = "qwertyuioplkjhgfdsazxcvbnm1234567890ZXCVBNMLKJHGFDSAQWERTYUIOP"
__StrBaseLen = len(__StrBase) - 1
__MaxImageSize = 3500 * 2000
__BorderWidth = 5
__AreaSizeLimit = 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 __distance(pa: "list", pb: "list") -> "float":
    return np.sqrt((pa[0] - pb[0]) ** 2 + (pa[1] - pb[1]) ** 2)


def __rect_area(points: "np.ndarray") -> "float":
    return __distance(points[0], points[1]) * __distance(points[0], points[3])


def __trans_img(img: "np.ndarray", points: "np.ndarray") -> "np.ndarray":
    tl, bl, br, tr = points
    width = int(max(__distance(tl, tr), __distance(bl, br)))
    height = int(max(__distance(tl, bl), __distance(tr, br)))
    rect = np.array([tl, tr, bl, br], dtype="float32")
    dest = np.array(
        [[0, 0], [width - 1, 0],
         [0, height - 1], [width - 1, height - 1]],
        dtype="float32"
    )

    matrix = cv2.getPerspectiveTransform(rect, dest)  # noqa
    return cv2.warpPerspective(img, matrix, (width, height))  # noqa


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(fmt: "str" = "%Y-%m-%d_%H-%M-%S") -> "str":
    return strftime(fmt, 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 resize_img(img: "np.ndarray", height: "int" = None, width: "int" = None) -> "np.ndarray":
    if height is None and width is None:
        return img
    h, w = img.shape[:2]
    if height is not None:
        dim = int(w * height / h), height
    else:
        dim = width, int(h * width / w)

    return cv2.resize(img, dim, interpolation=cv2.INTER_AREA)  # 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 preprocess_img(img: "np.ndarray") -> "np.ndarray":
    img = compress_img(img, 1000 * 600)
    # RGB -> gray
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # noqa
    # 二值化
    thresh_type = cv2.THRESH_BINARY  # noqa
    if __isWhiteBack(gray):
        thresh_type = cv2.THRESH_BINARY_INV  # noqa
    _, thresh = cv2.threshold(gray, 158, 255, thresh_type)  # noqa
    # 边缘检测
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # noqa
    if not contours:
        return img
    max_contour = max(contours, key=cv2.contourArea)  # noqa 找到最大的轮廓
    # 图像校正
    approx = cv2.approxPolyDP(max_contour, 0.02 * cv2.arcLength(max_contour, True), True)  # noqa
    area = img.shape[0] * img.shape[1]
    if len(approx) == 4:  # 较为规则，可以校正
        points = approx.reshape(4, 2)
        if __rect_area(points) * __AreaSizeLimit < area:
            return img
        return __trans_img(img, points)
    else:  # 不太规则，尝试裁剪
        rect = cv2.minAreaRect(max_contour)  # noqa 计算最小外接矩形
        box = np.intp(cv2.boxPoints(rect))  # noqa 获取矩形的四个角点
        if __rect_area(box) * __AreaSizeLimit < area:
            return img
        return img[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()