图像数据增强方法（opencv-python && PIL）

该脚本实现了一些数据图像增强的方法

import cv2
import numpy as np
import os
import random
from PIL import Image
from PIL import ImageEnhance
"""
该脚本实现了一些数据图像增强的方法
"""


# 高斯噪声
def gauss(image_path, save_path):

    # 设置高斯分布的均值和标准差
    mean_list = [-0.2, -0.1, 0, 0.3, 0.6]
    sigma_list = [40, 50, 60]

    files = os.listdir(image_path)

    for file in files:
        file_path = os.path.join(image_root, file)
        temp = os.path.splitext(file)
        file_name = temp[0] + '_gauss' + temp[-1]

        image = cv2.imread(file_path)
        mean = random.choice(mean_list)
        sigma = random.choice(sigma_list)
        gauss = np.random.normal(mean, sigma, image.shape)
        noisy_image = image + gauss

        # 将noisy_image中的像素控制在0-255
        noisy_image = np.clip(noisy_image, a_min=0, a_max=255)
        cv2.imwrite(os.path.join(save_path, file_name), noisy_image)


# 椒盐噪声
def s_p(image_path, save_path):

    # 设置椒盐噪声中 s p 的比例，以及椒盐所占的百分比
    s_p_list = [0.4, 0.5, 0.6]
    amount_list = [0.02, 0.03, 0.04]

    files = os.listdir(image_path)

    for file in files:

        file_path = os.path.join(image_root, file)
        temp = os.path.splitext(file)
        file_name = temp[0] + '_sp' + temp[-1]

        image = cv2.imread(file_path)
        noisy_image = np.copy(image)
        s_p = random.choice(s_p_list)
        amount = random.choice(amount_list)

        num_salt = np.ceil(amount * image.size * s_p)
        coords = [np.random.randint(0, i - 1, int(num_salt)) for i in image.shape]
        noisy_image[coords[0], coords[1], :] = [255, 255, 255]

        num_pepper = np.ceil(amount * image.size * (1. - s_p))
        coords = [np.random.randint(0, i - 1, int(num_pepper)) for i in image.shape]
        noisy_image[coords[0], coords[1], :] = [0, 0, 0]
        cv2.imwrite(os.path.join(save_path, file_name), noisy_image)


# 乘性噪声
def mult(image_path, save_path):

    # 设置乘性噪声的均值和sigma
    mean_list = [-0.2, -0.1, 0, 0.3, 0.6]
    sigma_list = [0.2, 0.4, 0.6, 0.8]

    files = os.listdir(image_path)

    for file in files:

        file_path = os.path.join(image_path, file)
        temp = os.path.splitext(file)
        file_name = temp[0] + '_speck' + temp[-1]

        image = cv2.imread(file_path)
        mean = random.choice(mean_list)
        sigma = random.choice(sigma_list)

        gauss = np.random.normal(mean, sigma, image.shape)
        noisy_image = image + image * gauss
        noisy_image = np.clip(noisy_image, a_min=0, a_max=255)
        cv2.imwrite(os.path.join(save_path, file_name), noisy_image)


# 图像高宽放大为1.5倍
def change_scale_15(image_path, save_path):

    files = os.listdir(image_path)

    for file in files:
        file_path = os.path.join(image_path, file)
        image = cv2.imread(file_path)
        image = cv2.resize(image, None, fx=1.5, fy=1.5)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_scale_1.5' + temp[-1]
        cv2.imwrite(os.path.join(save_path, file_name), image)


# 图像高宽缩小为0.5倍
def change_scale_05(image_path, save_path):

    files = os.listdir(image_path)

    for file in files:
        file_path = os.path.join(image_path, file)
        image = cv2.imread(file_path)
        image = cv2.resize(image, None, fx=0.5, fy=0.5)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_scale_0.5' + temp[-1]
        cv2.imwrite(os.path.join(save_path, file_name), image)


# 图像水平翻转
def horizontal_flip(image_path, save_path):

    files = os.listdir(image_path)

    for file in files:
        file_path = os.path.join(image_path, file)
        image = cv2.imread(file_path)
        image = cv2.flip(image, 1)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_horizontal' + temp[-1]
        cv2.imwrite(os.path.join(save_path, file_name), image)


# 图像垂直翻转
def vertical_flip(image_path, save_path):

    files = os.listdir(image_path)

    for file in files:
        file_path = os.path.join(image_path, file)
        image = cv2.imread(file_path)
        image = cv2.flip(image, 0)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_vertical' + temp[-1]
        cv2.imwrite(os.path.join(save_path, file_name), image)


# 图像顺时针或逆时针旋转90°
def rotate_90(image_path, save_path):

    files = os.listdir(image_path)
    angel_list = [0, 1]

    for file in files:
        file_path = os.path.join(image_path, file)
        image = cv2.imread(file_path)
        image = cv2.transpose(image)
        rotate_angle = random.choice(angel_list)
        image = cv2.flip(image, rotate_angle)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_90' + temp[-1]
        cv2.imwrite(os.path.join(save_path, file_name), image)


# 对比度增强
def contrast(image_path, save_path):

    files = os.listdir(image_path)

    for file in files:
        image = Image.open(os.path.join(image_path, file))
        image = ImageEnhance.Contrast(image)
        contrast_factor = np.random.randint(5, 19) / 10.

        image = image.enhance(contrast_factor)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_contrast' + temp[-1]
        image.save(os.path.join(save_path, file_name))
        # print('%s use %.1f' % (file_name, contrast))


# 亮度增强
def bright(image_path, save_path):

    files = os.listdir(image_path)

    for file in files:
        image = Image.open(os.path.join(image_path, file))
        image = ImageEnhance.Brightness(image)
        bright_factor = np.random.randint(5, 19) / 10.

        image = image.enhance(bright_factor)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_bright' + temp[-1]
        image.save(os.path.join(save_path, file_name))
        # print('%s use %.1f' % (file_name, bright_factor))


# 颜色增强
def color(image_path, save_path):

    files = os.listdir(image_path)

    for file in files:
        image = Image.open(os.path.join(image_path, file))
        image = ImageEnhance.Color(image)
        color_factor = np.random.randint(1, 25) / 10.

        image = image.enhance(color_factor)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_color' + temp[-1]
        image.save(os.path.join(save_path, file_name))
        # print('%s use %.1f' % (file_name, color_factor))


# 锐化增强
def sharp(image_path, save_path):

    files = os.listdir(image_path)

    for file in files:
        image = Image.open(os.path.join(image_path, file))
        image = ImageEnhance.Sharpness(image)
        sharp_factor = np.random.randint(1, 30) / 10.

        image = image.enhance(sharp_factor)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_sharp' + temp[-1]
        image.save(os.path.join(save_path, file_name))
        # print('%s use %.1f' % (file_name, sharp_factor))


# 混合增强
def hybrid(image_path, save_path):

    files = os.listdir(image_path)

    for file in files:
        image = Image.open(os.path.join(image_path, file))

        contrast_factor = np.random.randint(5, 19) / 10.
        bright_factor = np.random.randint(5, 19) / 10.
        color_factor = np.random.randint(1, 25) / 10.
        sharp_factor = np.random.randint(1, 30) / 10.

        image = ImageEnhance.Contrast(image).enhance(contrast_factor)
        image = ImageEnhance.Brightness(image).enhance(bright_factor)
        image = ImageEnhance.Color(image).enhance(color_factor)
        image = ImageEnhance.Sharpness(image).enhance(sharp_factor)

        temp = os.path.splitext(file)
        file_name = temp[0] + '_hybrid' + temp[-1]
        # cv2.imwrite(os.path.join(save_path, file_name), image)
        image.save(os.path.join(save_path, file_name))
        print('%s use contrast:%.1f bright:%.1f color:%.1f sharp:%.1f' %
              (file_name, contrast_factor, bright_factor, color_factor, sharp_factor))


image_root = r"C:\Users\Administrator\Desktop\test\photos"
save_root = r"C:\Users\Administrator\Desktop\test\new_photos"


# gauss(image_root, save_root)
# s_p(image_root, save_root)
# mult(image_root, save_root)
# change_scale_15(image_root, save_root)
# change_scale_05(image_root, save_root)
# horizontal_flip(image_root, save_root)
# vertical_flip(image_root, save_root)
# rotate_90(image_root, save_root)
# contrast(image_root, save_root)
# color(image_root, save_root)
# sharp(image_root, save_root)
# hybrid(image_root, save_root)

有用的话记得点赞关注一波！