VerificationCode/utils.py

import base64
import colorsys
import copy
import re
import socket
import traceback
import urllib
from enum import Enum
from functools import reduce

import requests
from PIL import ImageDraw, ImageFont, Image
from keras import backend as K
import numpy as np
import cv2


def compose(*funcs):
    """Compose arbitrarily many functions, evaluated left to right.
    Reference: https://mathieularose.com/function-composition-in-python/
    """
    if funcs:
        return reduce(lambda f, g: lambda *a, **kw: g(f(*a, **kw)), funcs)
    else:
        raise ValueError('Composition of empty sequence not supported.')


def box_iou(b1, b2):
    """Return iou tensor
    Parameters
    ----------
    b1: tensor, shape=(i1,...,iN, 4), xywh
    b2: tensor, shape=(j, 4), xywh
    Returns
    -------
    iou: tensor, shape=(i1,...,iN, j)
    """

    # Expand dim to apply broadcasting.
    b1 = K.expand_dims(b1, -2)
    b1_xy = b1[..., :2]
    b1_wh = b1[..., 2:4]
    b1_wh_half = b1_wh/2.
    b1_mins = b1_xy - b1_wh_half
    b1_maxes = b1_xy + b1_wh_half

    # Expand dim to apply broadcasting.
    b2 = K.expand_dims(b2, 0)
    b2_xy = b2[..., :2]
    b2_wh = b2[..., 2:4]
    b2_wh_half = b2_wh/2.
    b2_mins = b2_xy - b2_wh_half
    b2_maxes = b2_xy + b2_wh_half

    intersect_mins = K.maximum(b1_mins, b2_mins)
    intersect_maxes = K.minimum(b1_maxes, b2_maxes)
    intersect_wh = K.maximum(intersect_maxes - intersect_mins, 0.)
    intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
    b1_area = b1_wh[..., 0] * b1_wh[..., 1]
    b2_area = b2_wh[..., 0] * b2_wh[..., 1]
    iou = intersect_area / (b1_area + b2_area - intersect_area)

    return iou


def get_classes(classes_path):
    """loads the classes"""
    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names


def get_anchors(anchors_path):
    """loads the anchors from a file"""
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = [float(x) for x in anchors.split(',')]
    return np.array(anchors).reshape(-1, 2)


def get_colors(number, bright=True):
    """
    Generate random colors for drawing bounding boxes.
    To get visually distinct colors, generate them in HSV space then
    convert to RGB.
    """
    if number <= 0:
        return []

    brightness = 1.0 if bright else 0.7
    hsv_tuples = [(x / number, 1., brightness)
                  for x in range(number)]
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(
        map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
            colors))
    np.random.seed(10101)  # Fixed seed for consistent colors across runs.
    np.random.shuffle(colors)  # Shuffle colors to decorrelate adjacent classes.
    np.random.seed(None)  # Reset seed to default.
    return colors


labelType = Enum('labelType', ('LABEL_TOP_OUTSIDE',
                               'LABEL_BOTTOM_OUTSIDE',
                               'LABEL_TOP_INSIDE',
                               'LABEL_BOTTOM_INSIDE',))


def draw_boxes(image, out_boxes, out_classes, out_scores, class_names, colors):
    font = ImageFont.truetype(font='yolo_data/FiraMono-Medium.otf',
                              size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
    thickness = (image.size[0] + image.size[1]) // 300

    box_list = []
    for i, c in reversed(list(enumerate(out_classes))):
        predicted_class = class_names[c]
        box = out_boxes[i]
        score = out_scores[i]

        label = '{} {:.2f}'.format(predicted_class, score)
        draw = ImageDraw.Draw(image)
        label_size = draw.textsize(label, font)

        top, left, bottom, right = box
        top = max(0, np.floor(top + 0.5).astype('int32'))
        left = max(0, np.floor(left + 0.5).astype('int32'))
        bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
        right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
        # print(label, (left, top), (right, bottom))
        box_list.append([(left, top), (right, bottom)])

        if top - label_size[1] >= 0:
            text_origin = np.array([left, top - label_size[1]])
        else:
            text_origin = np.array([left, top + 1])

        # My kingdom for a good redistributable image drawing library.
        for i in range(thickness):
            draw.rectangle(
                [left + i, top + i, right - i, bottom - i],
                outline=colors[c])
        draw.rectangle(
            [tuple(text_origin), tuple(text_origin + label_size)],
            fill=colors[c])
        draw.text(text_origin, label, fill=(0, 0, 0), font=font)
        del draw

    return image, box_list


def draw_label(image, text, color, coords, label_type=labelType.LABEL_TOP_OUTSIDE):
    font = cv2.FONT_HERSHEY_PLAIN
    font_scale = 1.
    (text_width, text_height) = cv2.getTextSize(text, font, fontScale=font_scale, thickness=1)[0]

    padding = 5
    rect_height = text_height + padding * 2
    rect_width = text_width + padding * 2

    (x, y) = coords

    if label_type == labelType.LABEL_TOP_OUTSIDE or label_type == labelType.LABEL_BOTTOM_INSIDE:
        cv2.rectangle(image, (x, y), (x + rect_width, y - rect_height), color, cv2.FILLED)
        cv2.putText(image, text, (x + padding, y - text_height + padding), font,
                    fontScale=font_scale,
                    color=(255, 255, 255),
                    lineType=cv2.LINE_AA)
    else:
        # LABEL_BOTTOM_OUTSIDE or LABEL_TOP_INSIDE
        cv2.rectangle(image, (x, y), (x + rect_width, y + rect_height), color, cv2.FILLED)
        cv2.putText(image, text, (x + padding, y + text_height + padding), font,
                    fontScale=font_scale,
                    color=(255, 255, 255),
                    lineType=cv2.LINE_AA)

    return image


def pil_resize(image_np, height, width):
    image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
    image_pil = image_pil.resize((int(width), int(height)), Image.BICUBIC)
    image_np = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGB2BGR)
    return image_np


def pil_resize_a(image_np, height, width):
    image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGRA2RGBA))
    image_pil = image_pil.resize((int(width), int(height)), Image.BICUBIC)
    image_np = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGBA2BGRA)
    return image_np


def np2pil(image_np):
    image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
    return image_pil


def pil2np(image_pil):
    image_np = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    return image_np


def np2pil_a(image_np):
    image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGRA2RGBA))
    return image_pil


def pil2np_a(image_pil):
    image_np = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGBA2BGRA)
    return image_np


def pil_rotate(image_np, angle, fill_color=(255, 255, 255)):
    image_pil = np2pil(image_np)
    image_rotate = image_pil.rotate(angle, expand=False, fillcolor=fill_color)
    image_np = pil2np(image_rotate)
    return image_np


def pil_rotate_a(image_np, angle):
    image_pil = np2pil_a(image_np)
    image_rotate = image_pil.rotate(angle, expand=False, fillcolor=(255, 255, 255))
    image_np = pil2np_a(image_rotate)
    return image_np


def request_post(url, param, time_out=1000, use_zlib=False):
    fails = 0
    text = None
    while True:
        try:
            if fails >= 1:
                break

            headers = {'content-type': 'application/json'}
            # result = requests.post(url, data=param, timeout=time_out)
            session = requests.Session()
            result = session.post(url, data=param, timeout=time_out)

            if result.status_code == 200:
                text = result.text
                break
            else:
                print('result.status_code', result.status_code)
                print('result.text', result.text)
                fails += 1
                continue
        except socket.timeout:
            fails += 1
            print('timeout! fail times:', fails)
        except:
            fails += 1
            print('fail! fail times:', fails)
            traceback.print_exc()
    return text


def np2bytes(image_np):
    # numpy转为可序列化的string
    success, img_encode = cv2.imencode(".jpg", image_np)
    # numpy -> bytes
    img_bytes = img_encode.tobytes()
    return img_bytes


def bytes2np(_b, unchanged=False):
    try:
        # 二进制数据流转np.ndarray [np.uint8: 8位像素]
        if unchanged:
            # unchanged能读取透明通道
            image_np = cv2.imdecode(np.frombuffer(_b, np.uint8), cv2.IMREAD_UNCHANGED)
        else:
            image_np = cv2.imdecode(np.frombuffer(_b, np.uint8), cv2.IMREAD_COLOR)

        # 将rgb转为bgr
        # image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
        return image_np
    except cv2.error as e:
        if "src.empty()" in str(e):
            print("bytes2np image is empty!")
        return None
    except:
        traceback.print_exc()
        return None


def base64_decode(b64, is_test=False):
    # b64 = str(b64)
    # missing_padding = 4 - len(b64) % 4
    # if missing_padding:
    #     print('base64 missing_padding', missing_padding)
    #     b64 += '=' * missing_padding
    #     print('len(b64)', len(b64))

    if '%' in str(b64):
        b64 = urllib.parse.unquote(str(b64))
        print(b64)

    if type(b64) == bytes:
        return base64.b64decode(b64)
    b64 = re.sub(r"\r\n", "\n", b64)
    missing_padding = 4 - len(b64) % 4
    if missing_padding:
        b64 += '=' * missing_padding
    if is_test:
        return base64.b64decode(b64), b64
    return base64.b64decode(b64)


def limit_image_size(image_np, threshold=2000):
    h, w = image_np.shape[:2]
    if h <= threshold and w <= threshold:
        return image_np
    scale = threshold / max(h, w)
    h = int(h * scale)
    w = int(w * scale)
    image_np = pil_resize(image_np, h, w)
    return image_np


def get_best_predict_size(image_np, times=32):
    sizes = []
    for i in range(1, 100):
        if i*times <= 1300:
            sizes.append(i*times)
    sizes.sort(key=lambda x: x, reverse=True)

    min_len = 10000
    best_height = sizes[0]
    for height in sizes:
        if abs(image_np.shape[0] - height) < min_len:
            min_len = abs(image_np.shape[0] - height)
            best_height = height

    min_len = 10000
    best_width = sizes[0]
    for width in sizes:
        if abs(image_np.shape[1] - width) < min_len:
            min_len = abs(image_np.shape[1] - width)
            best_width = width

    return best_height, best_width


def add_contrast(image_np):
    # cv2.imshow("image_np", image_np)
    img = image_np.astype(np.float32)
    bri_mean = np.mean(img)

    # a = np.arange(5, 16, 5) / 10
    # b = np.arange(-30, 31, 30)
    #
    # a_len = len(a)
    # b_len = len(b)
    # print(a_len, b_len)
    #
    # for i in range(a_len):
    #     for j in range(b_len):
    #         aa = a[i]
    #         bb = b[j]
    #         img_a = aa * (img-bri_mean) + bb + bri_mean
    #         print(i, j, aa, bb)
    #         img_a = np.clip(img_a, 0, 255).astype(np.uint8)
    #         cv2.imshow("img_a", img_a)
    #         cv2.waitKey(0)

    aa = 3
    bb = -50

    img_a = aa * (img-bri_mean) + bb + bri_mean
    img_a = np.clip(img_a, 0, 255).astype(np.uint8)
    # cv2.imshow("img_a", img_a)
    # cv2.waitKey(0)
    return img_a


def image_to_str(image_np):
    file_bytes = np2bytes(image_np)
    file_base64 = base64.b64encode(file_bytes)
    file_str = file_base64.decode("utf-8")
    return file_str


def str_to_image(_str):
    b64 = _str.encode("utf-8")
    image_np = bytes2np(base64_decode(b64))
    return image_np


def rgba_to_rgb(image_np, return_position=False):
    # 提取 RGB 通道和透明通道
    rgb_channels = image_np[:, :, :3]
    alpha_channel = image_np[:, :, 3]

    # 找到透明通道小于250的部分
    transparent_pixels = alpha_channel < 250

    # 将透明通道小于250的部分的RGB值置为黑色
    rgb_channels[transparent_pixels] = [100, 100, 100]

    image_np = rgb_channels

    # 创建图像的副本，并将所有像素值置为白色
    # white_background = np.ones_like(image_np) * 255
    white_background = copy.copy(image_np)

    # 在图像副本上找到黑色部分的轮廓
    contours, hierarchy = cv2.findContours(transparent_pixels.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 输出黑色部分的位置（边界框坐标）
    position_list = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        position_list.append([int(x), int(y), int(x+w), int(y+h)])

    # 将黑色部分的轮廓描绘为白色
    cv2.drawContours(white_background, contours, -1, (255, 255, 255), thickness=1)

    # 合并原始图像和处理后的图像
    # image_np = cv2.addWeighted(image_np, 0, white_background, 1, 0)

    image_np = white_background

    if return_position:
        return image_np, position_list

    return image_np