utils/images.py

# -*- coding:utf-8 –*-
import random
from math import *

import cv2
from PIL import ImageDraw, Image, ImageFont

from utils.common import *
from utils.multiprogress import MultiThreading


def show(image, name="Press 'q' exit"):
    if not hasattr(image, 'shape'):
        img_array = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
    else:
        img_array = image.copy()
    cv2.imshow(name, img_array)
    cv2.moveWindow(name, 0, 0)
    key = cv2.waitKey(0) & 0xEFFFFF
    if key == ord('q'):
        cv2.destroyWindow(name)


def crop_face_square(img, box, pad=5):
    img_h, img_w, img_c = img.shape
    if box != [0, 0, 0, 0]:
        x0, y0, x1, y1 = box
        w, h = x1-x0, y1-y0
        # 扩充为正方形
        long_side = min(max(h, w) + int(2*pad), img_h)
        w_add, h_add = (long_side - w) // 2, (long_side - h) // 2
        crop_x0, crop_y0 = max(x0 - w_add, 0), max(y0 - h_add, 0)
        crop_x1, crop_y1 = min(crop_x0 + long_side, img_w), min(crop_y0 + long_side, img_h)
        crop_x0, crop_y0 = crop_x1 - long_side, crop_y1 - long_side
        return img[crop_y0:crop_y1, crop_x0:crop_x1], [crop_x0, crop_y0, crop_x1, crop_y1]
    else:
        # print('No detected box, crop right rect.')
        if img_h == 960:
            img = img[60:780, :]
            img_h, img_w = img.shape[:2]
        return img[:, img_w - img_h:img_w], [0, 0, 0, 0]


def crop_face_square_rate(img, box, rate=0.1):
    img_h, img_w, img_c = img.shape
    if box != [0, 0, 0, 0]:
        x0, y0, x1, y1 = box
        w, h = x1-x0, y1-y0
        # 扩充为正方形
        pad = max(w, h) * rate
        long_side = min(max(h, w) + int(2*pad), img_h)
        w_add, h_add = (long_side - w) // 2, (long_side - h) // 2
        crop_x0, crop_y0 = max(x0 - w_add, 0), max(y0 - h_add, 0)
        crop_x1, crop_y1 = min(crop_x0 + long_side, img_w), min(crop_y0 + long_side, img_h)
        crop_x0, crop_y0 = crop_x1 - long_side, crop_y1 - long_side
        return img[crop_y0:crop_y1, crop_x0:crop_x1], [crop_x0, crop_y0, crop_x1, crop_y1]
    else:
        if img_h == 960:
            crop_x0, crop_x1 = img_w - 720, img_w
            crop_y0, crop_y1 = 60, 780
        else:
            crop_x0, crop_x1 = img_w - img_h, img_w
            crop_y0, crop_y1 = 0, img_h
        return img[crop_y0:crop_y1, crop_x0:crop_x1], [crop_x0, crop_y0, crop_x1, crop_y1]


def expand_box_rate(img, box, rate=0.1):
    img_h, img_w, img_c = img.shape
    if box != [0, 0, 0, 0]:
        x0, y0, x1, y1 = box
        w, h = x1-x0, y1-y0
        # 扩充为正方形
        pad = max(w, h) * rate
        long_side = min(max(h, w) + int(2*pad), img_h)
        w_add, h_add = (long_side - w) // 2, (long_side - h) // 2
        crop_x0, crop_y0 = x0 - w_add, y0 - h_add
        crop_x1, crop_y1 = crop_x0 + long_side, crop_y0 + long_side
        return [crop_x0, crop_y0, crop_x1, crop_y1]
    else:
        if img_h == 960:
            crop_x0, crop_x1 = img_w - 720, img_w
            crop_y0, crop_y1 = 60, 780
        else:
            crop_x0, crop_x1 = img_w - img_h, img_w
            crop_y0, crop_y1 = 0, img_h
        return [crop_x0, crop_y0, crop_x1, crop_y1]


def crop_with_pad(img, crop_box):
    img_h, img_w, img_c = img.shape
    x0, y0, x1, y1 = crop_box
    w, h = x1 - x0, y1 - y0
    if tuple(crop_box) == (0, 0, 0, 0) or w <= 50 or h <= 50:  # 背景，裁右半图
        if img_h == 960:
            crop_x0, crop_x1 = img_w - 720, img_w
            crop_y0, crop_y1 = 60, 780
        else:
            crop_x0, crop_x1 = img_w - img_h, img_w
            crop_y0, crop_y1 = 0, img_h
        crop_img = img[crop_y0:crop_y1, crop_x0:crop_x1]
        return crop_img
    else:
        crop_x0, crop_y0 = max(x0, 0), max(y0, 0)
        crop_x1, crop_y1 = min(x1, img_w), min(y1, img_h)
        left, top, right, bottom = crop_x0 - x0, crop_y0 - y0, x1 - crop_x1, y1 - crop_y1
        crop_img = img[crop_y0:crop_y1, crop_x0:crop_x1]
        crop_img = cv2.copyMakeBorder(crop_img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(128, 128, 128))
        return crop_img


def clip_paste_rect(bg_img, fg_img, loc, refine=True):
    """
    选取粘贴区域，不要超出背景范围
    loc: 中心点(cx, cy)
    """
    bg_h, bg_w = bg_img.shape[:2]
    fg_h, fg_w = fg_img.shape[:2]
    fg_h = fg_h - 1 if fg_h % 2 else fg_h
    fg_w = fg_w - 1 if fg_w % 2 else fg_w
    cx, cy = loc
    left, top, right, bottom = cx - fg_w // 2, cy - fg_h // 2, cx + fg_w // 2, cy + fg_h // 2

    if refine:
        right, bottom = min(right, bg_w), min(bottom, bg_h)
        left, top = right - fg_w, bottom - fg_h
        left, top = max(0, left), max(0, top)
        right, bottom = left + fg_w, top + fg_h

    plot_x1, plot_y1, plot_x2, plot_y2 = left, top, right, bottom
    use_x1, use_y1, use_x2, use_y2 = 0, 0, fg_w, fg_h

    if left < 0:
        plot_x1, use_x1 = 0, -left
    if top < 0:
        plot_y1, use_y1 = 0, -top
    if right > bg_w:
        plot_x2, use_x2 = bg_w, fg_w - (right - bg_w)
    if bottom > bg_h:
        plot_y2, use_y2 = bg_h, fg_h - (bottom - bg_h)

    use_bg = bg_img[plot_y1:plot_y2, plot_x1:plot_x2]
    use_fg = fg_img[use_y1:use_y2, use_x1:use_x2]
    window = (plot_x1, plot_y1, plot_x2, plot_y2)

    return use_bg, use_fg, window


# @clock_custom(fmt=DEFAULT_FMT)
def paste(bg_img, fg_img, loc, trans_thresh=1, refine=True):
    """
    贴图
    loc: center (cx, cy)
    """
    use_bg, use_fg, window = clip_paste_rect(bg_img, fg_img, loc, refine)
    plot_x1, plot_y1, plot_x2, plot_y2 = window
    b, g, r, a = cv2.split(use_fg)
    a[a > 0] = 255
    a = np.dstack([a, a, a]) * trans_thresh
    use_bg = use_bg * (255.0 - a) / 255
    use_bg += use_fg[:, :, :3] * (a / 255)
    bg_img[plot_y1:plot_y2, plot_x1:plot_x2] = use_bg.astype('uint8')
    return bg_img


def put_chinese_text(img, text, position, font, text_color=(0, 255, 0)):
    if isinstance(img, np.ndarray):  # 判断是否OpenCV图片类型
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(img)
    draw.text(position, text, text_color, font=font)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)


def compute_mean_std(img_list, preprocess=None, workers=12):
    """
    计算公式：
    S^2 = sum((x - x')^2) / N = sum(x^2 + x'^2 - 2xx') / N
        = (sum(x^2) + sum(x'^2) - 2x'*sum(x)) / N
        = (sum(x^2) + N*(x'^2) - 2x'*(N * x')) / N
        = (sum(x^2) - N * (x'^2)) / N
        = sum(x^2) / N - (sum(x) / N)^2
        = mean(x^2) - mean(x)^2 = E(x^2) - E(x)^2
    :param img_list:
    :param workers:
    :param preprocess: 图像预处理函数，需要有BGR->RGB以及归一化操作
    :return: RGB means, stds
    """

    def cal_func(info):
        i, img_path = info
        img = cv2.imread(img_path)
        if preprocess:
            img = preprocess(img)
        else:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = img.astype(np.float32) / 255
        img = img.reshape((-1, 3))
        channel_mean = np.mean(img, axis=0)
        channel_var = np.mean(img**2, axis=0)
        if i % 1000 == 0:
            print(f"{i}/{len(img_list)}")
        return channel_mean, channel_var

    exe = MultiThreading(list(enumerate(img_list)), workers=min(workers, len(img_list)))
    res = exe.run(cal_func)
    all_means = np.array([r[0] for r in res])
    all_vars = np.array([r[1] for r in res])
    means = np.mean(all_means, axis=0)
    vars = np.mean(all_vars, axis=0)
    stds = np.sqrt(vars - means ** 2)

    return means, stds


def draw_province(f, val):
    img = Image.new("RGB", (45, 70), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    draw.text((0, 3), val, (0, 0, 0), font=f)
    img = img.resize((23, 70))
    char = np.array(img)
    return char


def draw_chars(f, val):
    img = Image.new("RGB", (23, 70), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    draw.text((0, 2), val, (0, 0, 0), font=f)
    A = np.array(img)
    return A


# 双层车牌
def draw_province_double(f, val):
    img = Image.new("RGB", (60, 60), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    draw.text((0, -12), val, (0, 0, 0), font=f)
    img = img.resize((80, 60))
    # img.show()
    char = np.array(img)
    return char


def draw_chars_ceil(f, val):
    img = Image.new("RGB", (30, 45), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    draw.text((1, -12), val, (0, 0, 0), font=f)
    img = img.resize((80, 60))
    # img.show()
    char = np.array(img)
    return char


def draw_chars_floor(f, val):
    img = Image.new("RGB", (30, 45), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    draw.text((1, -12), val, (0, 0, 0), font=f)
    img = img.resize((65, 110))
    # img.show()
    char = np.array(img)
    return char


# 图片做旧处理
def add_smudginess(img):
    smu = cv2.imread('data/bgp/smu.jpg')
    img_h, img_w = img.shape[:2]
    rows = r(smu.shape[0] - img_h)
    cols = r(smu.shape[1] - img_w)
    adder = smu[rows:rows + img_h, cols:cols + img_w]
    adder = cv2.resize(adder, (img_w, img_h))
    adder = cv2.bitwise_not(adder)
    val = random.random() * 0.5
    img = cv2.addWeighted(img, 1 - val, adder, val, 0.0)
    return img


def rot(img, angel, shape, max_angel):
    """ 使图像轻微的畸变
        img 输入图像
        factor 畸变的参数
        size 为图片的目标尺寸
    """
    size_o = [shape[1], shape[0]]
    size = (shape[1] + int(shape[0] * sin((float(max_angel)/180) * 3.14)), shape[0])
    interval = abs(int(sin((float(angel) / 180) * 3.14) * shape[0]))
    pts1 = np.float32([[0, 0], [0, size_o[1]], [size_o[0], 0], [size_o[0], size_o[1]]])
    if angel > 0:
        pts2 = np.float32([[interval, 0], [0, size[1]], [size[0], 0], [size[0]-interval, size_o[1]]])
    else:
        pts2 = np.float32([[0, 0], [interval, size[1]], [size[0]-interval, 0], [size[0], size_o[1]]])
    M = cv2.getPerspectiveTransform(pts1, pts2)
    dst = cv2.warpPerspective(img, M, size)
    return dst


def random_rot(img, factor, size):
    shape = size
    pts1 = np.float32([[0, 0], [0, shape[0]], [shape[1], 0], [shape[1], shape[0]]])
    pts2 = np.float32([[r(factor), r(factor)],
                        [r(factor), shape[0] - r(factor)],
                        [shape[1] - r(factor), r(factor)],
                        [shape[1] - r(factor), shape[0] - r(factor)]])
    M = cv2.getPerspectiveTransform(pts1, pts2)
    dst = cv2.warpPerspective(img, M, size)
    return dst


def random_rot_expend(img, factor, size):
    height, width = size
    degree = factor
    # 旋转后的尺寸
    n_h = int(width * fabs(sin(radians(degree))) + height * fabs(cos(radians(degree))))
    n_w = int(height * fabs(sin(radians(degree))) + width * fabs(cos(radians(degree))))
    M = cv2.getRotationMatrix2D((width / 2, height / 2), degree, 1)
    M[0, 2] += (n_w - width) / 2  # 重点在这步，目前不懂为什么加这步
    M[1, 2] += (n_h - height) / 2  # 重点在这步
    rot_img = cv2.warpAffine(img, M, (n_w, n_h), borderValue=(0, 0, 0))
    return rot_img


def random_rot_keep_size_left_right(img, factor1, factor2, size):
    # 透视变换 factor1 大于0 factor2 可正可负
    shape = size
    width = shape[0]
    height = shape[1]
    pts1 = np.float32([[0, 0], [width - 1, 0], [0, height-1], [width - 1, height-1]])
    point1_x = 0
    point1_y = 0
    point2_x = width - r(factor1)
    point2_y = r(factor2)  # shape[0] - r(factor)
    point3_x = 0  # shape[1] - r(factor)
    point3_y = height  # r(factor)
    point4_x = width - r(factor1)  # shape[1]
    point4_y = height - r(factor2)  # shape[0]
    max_x = max(point2_x, point4_x)
    max_y = point3_y
    if factor2 < 0:
        point1_x = 0
        point1_y = 0 - r(factor2)
        point2_x = width - r(factor1)
        point2_y = 0
        point3_x = 0
        point3_y = height + r(factor2)
        point4_x = width - r(factor1)
        point4_y = height
        max_x = max(point2_x, point4_x)
        max_y = point4_y
    pts2 = np.float32([[point1_x, point1_y],
                       [point2_x, point2_y],
                       [point3_x, point3_y],
                       [point4_x, point4_y]])
    M = cv2.getPerspectiveTransform(pts1, pts2)
    size2 = (max_x, max_y)
    dst = cv2.warpPerspective(img, M, size2)  # cv2.warpPerspective(img, M, size)
    return dst


# 仿射变换
def affine_transform(img, factor, size):
    shape = size
    pts1 = np.float32([[0, shape[0]], [shape[1], 0], [shape[1], shape[0]]])
    pts2 = np.float32([[r(factor), shape[0] - r(factor)],
                       [shape[1] - r(factor), r(factor)],
                       [shape[1], shape[0]]])  # [shape[1] - r(factor), shape[0] - r(factor)]])
    M = cv2.getAffineTransform(pts1, pts2)
    dst = cv2.warpAffine(img, M, size)
    return dst


# 腐蚀
def cv_erode(img, factor):
    value = r(factor)+1
    kernel = np.ones((value, value), np.uint8) * factor
    erosion = cv2.erode(img, kernel, iterations=1)
    return erosion


# 膨胀
def cv_dilate(img, factor):
    value = r(factor)+1
    kernel = np.ones((value, value), np.uint8)
    dilate = cv2.dilate(img, kernel, iterations=1)
    return dilate


def add_random_noise(img, factor):
    value = r(factor)
    for k in range(value):  # Create 5000 noisy pixels
        i = random.randint(0, img.shape[0] - 1)
        j = random.randint(0, img.shape[1] - 1)
        color = (random.randrange(256), random.randrange(256), random.randrange(256))
        img[i, j] = color
    return img


# 生成卷积核和锚点
def gen_kernel_anchor(length, angle):
    half = length / 2
    EPS = np.finfo(float).eps
    alpha = (angle - floor(angle / 180) * 180) / 180 * pi
    cosalpha = cos(alpha)
    sinalpha = sin(alpha)
    if cosalpha < 0:
        xsign = -1
    elif angle == 90:
        xsign = 0
    else:
        xsign = 1
    psfwdt = 1
    # 模糊核大小
    sx = int(fabs(length * cosalpha + psfwdt * xsign - length * EPS))
    sy = int(fabs(length * sinalpha + psfwdt - length * EPS))
    psf1 = np.zeros((sy, sx))
    # psf1是左上角的权值较大，越往右下角权值越小的核。
    # 这时运动像是从右下角到左上角移动
    for i in range(0, sy):
        for j in range(0, sx):
            psf1[i][j] = i * fabs(cosalpha) - j * sinalpha
            rad = sqrt(i * i + j * j)
            if rad >= half and fabs(psf1[i][j]) <= psfwdt:
                temp = half - fabs((j + psf1[i][j] * sinalpha) / cosalpha)
                psf1[i][j] = sqrt(psf1[i][j] * psf1[i][j] + temp * temp)
            psf1[i][j] = psfwdt + EPS - fabs(psf1[i][j])
            if psf1[i][j] < 0:
                psf1[i][j] = 0
    # 运动方向是往左上运动，锚点在（0，0）
    anchor = (0, 0)
    # 运动方向是往右上角移动，锚点一个在右上角
    # 同时，左右翻转核函数，使得越靠近锚点，权值越大
    if 0 < angle < 90:
        psf1 = np.fliplr(psf1)
        anchor = (psf1.shape[1] - 1, 0)
    elif -90 < angle < 0:  # 同理：往右下角移动
        psf1 = np.flipud(psf1)
        psf1 = np.fliplr(psf1)
        anchor = (psf1.shape[1] - 1, psf1.shape[0] - 1)
    elif angle < -90:  # 同理：往左下角移动
        psf1 = np.flipud(psf1)
        anchor = (0, psf1.shape[0] - 1)
    psf1 = psf1 / psf1.sum()
    return psf1, anchor


def hsv_transform(img):
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    # hsv[:, :, 0] = hsv[:, :, 0] * (0.2 + np.random.random() * 0.8)
    hsv[:, :, 1] = hsv[:, :, 1] * (0.5 + np.random.random()*0.5)
    hsv[:, :, 2] = hsv[:, :, 2] * (0.5 + np.random.random()*0.5)
    img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    return img


def kth_diag_indices(array, k=0, n=-1):
    """
    第k个对角线索引
    :param array: 输入二维矩阵
    :param k: k=0 主对角线 k>0 右移 k<0 左移
    :param n: n=-1, 对角线所有元素，否则只返回长为n的列表
    :return: 索引
    """
    if n == -1:
        rows, cols = np.diag_indices_from(array)
    else:
        rows, cols = np.diag_indices(n)
    if k < 0:
        return rows[-k:], cols[:k]
    elif k > 0:
        return rows[:-k], cols[k:]
    else:
        return rows, cols


def slash_mask(mask, start_index, end_index, set_value=0.0, mod='l', length=-1):
    """
    制作斜条纹的蒙版
    :param mask:
    :param start_index:
    :param end_index:
    :param set_value:
    :param mod: 左斜还是右斜
    :param length: 长度
    :return:
    """
    h, w = mask.shape[:2]
    assert length <= min(h, w)
    if mod == 'r':
        mask = np.fliplr(mask)
    for i in range(start_index, end_index+1):
        mask[kth_diag_indices(mask, i, length)] = set_value
    if mod == 'r':
        mask = np.fliplr(mask)
    return mask


def line_mask(mask, start_index, end_index, set_value=0.0, mod='h', length=-1):
    """
    制作条纹蒙版
    :param mask:
    :param start_index:
    :param end_index:
    :param set_value:
    :param mod: h 横 v 竖
    :param length: 长度
    :return:
    """
    h, w = mask.shape[:2]
    if mod == 'h':
        assert length <= w
        assert 0 <= start_index < end_index < h
        if length == -1 or length == w:
            mask[start_index: end_index+1, :] = set_value
        else:
            left = random.randint(0, w-length-1)
            right = left + length
            mask[start_index: end_index + 1, left:right] = set_value
    else:
        assert length <= h
        assert 0 <= start_index < end_index < w
        if length == -1 or length == h:
            mask[:, start_index: end_index + 1] = set_value
        else:
            top = random.randint(0, h - length - 1)
            bottom = top + length
            mask[top:bottom + 1, start_index: end_index] = set_value
    return mask


def read_binary_images(file_path):
    img = np.fromfile(file_path, dtype=np.uint8)
    img = np.reshape(img, (720, 1280, 3))
    cv2.imshow('img', img)
    cv2.waitKey(0)
    print()



if __name__ == '__main__':
    read_binary_images('/Users/didi/Desktop/座次/pic/bgr1641010923242.bgr')