python-OpenCV图片增强

import numpy as np
import cv2


def crop_image(img, x0, y0, w, h):
    """
    定义裁剪函数
    :param img: 要处理的图片
    :param x0: 左上角横坐标
    :param y0: 左上角纵坐标
    :param w: 裁剪宽度
    :param h: 裁剪高度
    :return: 裁剪后的图片
    """
    return img[x0:x0+w, y0:y0+h]


def random_crop(img, area_ratio, hw_vari):
    """
    定义随机裁剪函数
    :param img: 要处理的图片
    :param area_ratio: 裁剪画面占原图片的比例
    :param hw_vari: 扰动占原宽高的比例
    :return: 裁剪后的图片
    """
    w, h = img.shape[:2]
    hw_delta = np.random.uniform(-hw_vari, hw_vari)
    hw_mult = 1 + hw_delta

    # 裁剪宽度
    w_crop = int(np.round(w * np.sqrt(area_ratio * hw_mult)))
    # 裁剪宽度不能大于原图宽度
    if w_crop > w:
        w_crop = w
    # 裁剪高度
    h_crop = int(np.round(h * np.sqrt(area_ratio / hw_mult)))
    if h_crop > h:
        h_crop = h

    # 随机生成左上角的位置
    x0 = np.random.randint(0, w - w_crop + 1)
    y0 = np.random.randint(0, h - h_crop + 1)

    return crop_image(img, x0, y0, w_crop, h_crop)


def rotate_image(img, angle, crop):
    """
    定义旋转函数
    :param img: 要处理的图片
    :param angle: 图片旋转的角度
    :param crop: 是否对旋转后出现的黑边进行裁剪
    :return: 旋转后的图片
    """
    w, h = img.shape[:2]

    # 旋转角度的周期是360°
    angle %= 360

    # 计算仿射变换矩阵
    M_rotate = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1)
    # 得到旋转后的图像
    img_rotated = cv2.warpAffine(img, M_rotate, (w, h))

    # 如果需要裁剪黑边
    if crop:
        # 裁剪角度的等效周期是180°
        angle_crop = angle % 180
        if angle_crop > 90:
            angle_crop = 180 - angle_crop

        # 角度转化为弧度
        theta = angle_crop * np.pi / 180

        # 计算高宽比
        hw_ratio = float(h) / float(w)

        # 计算裁剪边长系数的分子项
        tan_theta = np.tan(theta)
        numerator = np.cos(theta) + np.sin(theta) * tan_theta

        # 计算分母项中相关的项
        r = hw_ratio if h > w else 1 / hw_ratio
        # 分母项
        denominator = r * tan_theta + 1

        # 最终的边长系数
        crop_mult = numerator / denominator

        # 得到裁剪区域
        w_crop = int(round(crop_mult * w))
        h_crop = int(round(crop_mult * h))
        x0 = int((w - w_crop) / 2)
        y0 = int((h - h_crop) / 2)

        img_rotated = crop_image(img_rotated, x0, y0, w_crop, h_crop)

    return img_rotated


def random_rotate(img, angle_vari, p_crop):
    """
    定义随机旋转函数
    :param img: 要处理的图片
    :param angle_vari: 旋转角度的范围
    :param p_crop: 要进行裁剪黑边的图片所占的比例
    :return: 随机旋转的图片
    """
    angle = np.random.uniform(-angle_vari, angle_vari)
    crop = False if np.random.random() > p_crop else True

    return rotate_image(img, angle, crop)


def hsv_transform(img, hue_delta, sat_mult, val_mult):
    """
    定义hsv转换函数
    :param img: 要处理的图片
    :param hue_delta: 色调变化
    :param sat_mult: 饱和度变化
    :param val_mult: 明暗度变化
    :return: hsv变化后的图片
    """
    img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float)
    img_hsv[:, :, 0] = (img_hsv[:, :, 0] + hue_delta) % 180
    img_hsv[:, :, 1] *= sat_mult
    img_hsv[:, :, 2] *= val_mult
    img_hsv[img_hsv > 255] = 255

    return cv2.cvtColor(np.round(img_hsv).astype(np.uint8), cv2.COLOR_HSV2BGR)


def random_hsv_transform(img, hue_vari, sat_vari, val_vari):
    """
    随机hsv变换
    :param img: 要处理的图片
    :param hue_vari: 色调比例变化范围
    :param sat_vari: 饱和度比例变化范围
    :param val_vari: 明暗度比例变化范围
    :return: 随机hsv变换后的图片
    """
    hue_delta = np.random.uniform(-hue_vari, hue_vari)
    sat_mult = np.random.uniform(-sat_vari, sat_vari) + 1
    val_mult = np.random.uniform(-val_vari, val_vari) + 1

    return hsv_transform(img, hue_delta, sat_mult, val_mult)


def gamma_transform(img, gamma):
    """
    定义gamma变换函数
    :param img: 要处理的函数
    :param gamma: gamma系数
    :return: gamma变换后的图片
    """
    gamma_table = [np.power(x / 255, gamma) * 255 for x in range(256)]
    gamma_table = np.round(np.array(gamma_table)).astype(np.uint8)

    return cv2.LUT(img, gamma_table)


def random_gamma_transform(img, gamma_vari):
    log_gamma_vari = np.log(gamma_vari)
    alpha = np.random.uniform(-log_gamma_vari, log_gamma_vari)
    gamma = np.exp(alpha)

    return gamma_transform(img, gamma)