OpenCV—python 图像增强

版权声明：转载请说明来源，谢谢 https://blog.csdn.net/wsp_1138886114/article/details/83184041

一、图像增强（该代码有问题）

import os
import cv2
import random
import numpy as np


def random_crop(img, area_ratio, hw_vari):
    """
    :param img:
    :param area_ratio: 裁剪画面占原画面比例(0,0.5)
    :param hw_vari:    扰动占原高宽比的范围(0.2)
    x0: 左上角横坐标
    y0: 左上角纵坐标
    w:  裁剪宽度
    h:  裁剪高度
    :return:
    """
    h, w = img.shape[:2]
    hw_delta = np.random.uniform(-hw_vari, hw_vari)
    hw_mult = 1 + hw_delta
    w_crop = int(round(w*np.sqrt(area_ratio*hw_mult)))
    if w_crop > w - 2:
        w_crop = w - 2
    h_crop = int(round(h*np.sqrt(area_ratio/hw_mult)))
    if h_crop > h - 2:
        h_crop = h - 2
    x0 = np.random.randint(0, w-w_crop-1)
    y0 = np.random.randint(0, h-h_crop-1)
    crop_image = lambda img, x0, y0, w, h: img[y0:y0 + h, x0:x0 + w]
    return crop_image


def random_rotate(img, angle_vari, p_crop):
    """
    :param img:
    :param angle_vari: 旋转角度（逆时针）180
    :param p_crop:     去黑边裁剪比例0.5
    :return:
    """
    angle = np.random.uniform(-angle_vari, angle_vari)
    crop = False if np.random.random() > p_crop else True

    h, w = img.shape[:2]
    angle %= 360
    M_rotate = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1)
    img_rotated = cv2.warpAffine(img, M_rotate, (w, h))

    if crop:
        angle_crop = angle % 180
        if angle_crop > 90:
            angle_crop = 180 - angle_crop
        theta = angle_crop * np.pi / 180.0
        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 = lambda img, x0, y0, w, h: img[y0:y0 + h, x0:x0 + w]
    return img_rotated


def random_hsv_transform(img, hue_vari, sat_vari, val_vari):
    """
    :param img:
    :param hue_vari: 色调变化比例范围
    :param sat_vari: 饱和度变化比例范围
    :param val_vari: 明度变化比例范围
    :return:
    """
    hue_delta = np.random.randint(-hue_vari, hue_vari)
    sat_mult = 1 + np.random.uniform(-sat_vari, sat_vari)
    val_mult = 1 + np.random.uniform(-val_vari, val_vari)

    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_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)
    gamma_table = [np.power(x / 255.0, gamma) * 255.0 for x in range(256)]
    gamma_table = np.round(np.array(gamma_table)).astype(np.uint8)
    return cv2.LUT(img, gamma_table)

def random_flip_img(img):
    """
    0 = X axis, 1 = Y axis,  -1 = both
    :param img:
    :return:
    """
    flip_val = [0,1,-1]
    random_flip_val = random.choice(flip_val)
    res = cv2.flip(img, random_flip_val)
    return res

if __name__ == '__main__':

    output_dir = "./temp_dir"
    input_dir = "./temp_dir"
    index = 1
    for child_dir in os.listdir(input_dir):
        child_path = os.path.join(input_dir, child_dir)
        for dir_image in os.listdir(child_path):
            print(child_path)
            Suffix_name = ['.png', '.jpg', '.jpeg']
            if dir_image.endswith(tuple(Suffix_name)):
                img = cv2.imread(os.path.join(child_path, dir_image))

                random_crop_img = random_crop(img,0.2,0.2)
                cv2.imwrite(child_path + '/crop_img' + str(index) + '.png', random_crop_img)

                res_rotate = random_rotate(img, 180, 0.5)
                cv2.imwrite(child_path + '/rotate_img' + str(index) + '.png', res_rotate)

                HSV_IMG = random_hsv_transform(img,2,2,3)
                cv2.imwrite(child_path + '/HSV_IMG' + str(index) + '.png', HSV_IMG)

                GAMMA_IMG = random_gamma_transform(img,0.2)
                cv2.imwrite(child_path + '/GAMMA_IMG' + str(index) + '.png', GAMMA_IMG)

                res_flip = random_flip_img(img)
                cv2.imwrite(child_path + './flip_img'+ str(index) + '.png', res_flip)

                index+=1

二、图像增强

获取目录下所有文件夹以及其里面的图片文件并对其执行图像增强

import os
import cv2
import random
import numpy as np


def Affine_transformation(img):
    """
    :param img:
    :return:
    """
    h, w = img.shape[:2]
    points1 = np.float32([[50, 50], [200, 50], [50, 200]])
    points2 = np.float32([[10, 100], [200, 50], [100, 250]])

    matrix = cv2.getAffineTransform(points1, points2)
    Affine_transfor_img = cv2.warpAffine(img, matrix, (h, w))
    return Affine_transfor_img

def random_rotate_img(img, min_angle, max_angle):
    '''
        random rotation an image

    :param img:         image to be rotated
    :param min_angle:   min angle to rotate
    :param max_angle:   max angle to rotate
    :return:            image after random rotated
    '''
    if not isinstance(img, list):
        img = [img]

    angle = random.randint(min_angle, max_angle)
    center = (img[0].shape[0] / 2, img[0].shape[1] / 2)
    rot_matrix = cv2.getRotationMatrix2D(center, angle, scale=1.0)

    res = []
    for img_inst in img:
        img_inst = cv2.warpAffine(img_inst, rot_matrix, dsize=img_inst.shape[:2], borderMode=cv2.BORDER_CONSTANT)
        res.append(img_inst)
    if len(res) == 0:
        res = res[0]
    return res

def random_hsv_transform(img, hue_vari, sat_vari, val_vari):
    """
    :param img:
    :param hue_vari: 色调变化比例范围(0,360)
    :param sat_vari: 饱和度变化比例范围(0,1)
    :param val_vari: 明度变化比例范围(0,1)
    :return:
    """
    hue_delta = np.random.randint(-hue_vari, hue_vari)
    sat_mult = 1 + np.random.uniform(-sat_vari, sat_vari)
    val_mult = 1 + np.random.uniform(-val_vari, val_vari)

    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_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)
    gamma_table = [np.power(x / 255.0, gamma) * 255.0 for x in range(256)]
    gamma_table = np.round(np.array(gamma_table)).astype(np.uint8)
    return cv2.LUT(img, gamma_table)

def random_flip_img(img):
    """
    0 = X axis, 1 = Y axis,  -1 = both
    :param img:
    :return:
    """
    flip_val = [0,1,-1]
    random_flip_val = random.choice(flip_val)
    res = cv2.flip(img, random_flip_val)
    return res

if __name__ == '__main__':

    output_dir = "./face_dataset/test_faces"
    input_dir = "./face_dataset/test_faces"
    index = 1
    for child_dir in os.listdir(input_dir):
        child_path = input_dir +"/"+ child_dir
        for dir_image in os.listdir(child_path):
            Suffix_name = ['.png', '.jpg', '.jpeg']
            if dir_image.endswith(tuple(Suffix_name)):
                img = cv2.imread(child_path +"/"+ dir_image)

                Affine_transfor_img = Affine_transformation(img)
                cv2.imwrite(child_path + "/" + str(dir_image[:-4]) + '_Affine_transfor_img' + str(index) + '.png', Affine_transfor_img)

                res_rotate = random_rotate_img(img, 20, 70)
                cv2.imwrite(child_path +"/" +str(dir_image[:-4])+ '_rotate_img' + str(index) + '.png', res_rotate[0])

                HSV_IMG = random_hsv_transform(img,2,0.3,0.6)
                cv2.imwrite(child_path +"/" +str(dir_image[:-4])+'_HSV_IMG' + str(index) + '.png', HSV_IMG)

                GAMMA_IMG = random_gamma_transform(img,0.3)
                cv2.imwrite(child_path +"/" +str(dir_image[:-4])+'_GAMMA_IMG' + str(index) + '.png', GAMMA_IMG)

                res_flip = random_flip_img(img)
                cv2.imwrite(child_path +"/" +str(dir_image[:-4])+'_flip_img'+ str(index) + '.png', res_flip)
                index+=1