数据集增强——图像翻转、旋转、随机颜色、对比度、亮度、颜色增强

本代码共采用了四种数据增强，如采用其他数据增强方式，可以参考本代码，随意替换。

'''
imageDir 为原数据集的存放位置
saveDir  为数据增强后数据的存放位置
'''
 
from PIL import Image
from PIL import ImageEnhance
import os
import cv2
import numpy as np
import glob
 
def flip(root_path):   #翻转图像
    img = Image.open(root_path)
    filp_img = img.transpose(Image.FLIP_LEFT_RIGHT)
    return filp_img
 
def rotation(root_path):
    img = Image.open(root_path)
    rotation_img = img.rotate(90) #旋转角度
    return rotation_img
 
def brightnessEnhancement(root_path):#亮度增强
    image = Image.open(root_path)
    enh_bri = ImageEnhance.Brightness(image)
    brightness = 1.5
    image_brightened = enh_bri.enhance(brightness)
    return image_brightened
 
def contrastEnhancement(root_path):  # 对比度增强
    image = Image.open(root_path)
    enh_con = ImageEnhance.Contrast(image)
    contrast = 1.5
    image_contrasted = enh_con.enhance(contrast)
    return image_contrasted
 
def colorEnhancement(root_path):#颜色增强
    image = Image.open(root_path)
    enh_col = ImageEnhance.Color(image)
    color = 1.5
    image_colored = enh_col.enhance(color)
    return image_colored
 
def randomColor(root_path): #随机颜色
    """
    对图像进行颜色抖动
    :param image: PIL的图像image
    :return: 有颜色色差的图像image
    """
    image = Image.open(root_path)
    random_factor = np.random.randint(0, 31) / 10.  # 随机因子
    color_image = ImageEnhance.Color(image).enhance(random_factor)  # 调整图像的饱和度
    random_factor = np.random.randint(10, 21) / 10.  # 随机因子
    brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor)  # 调整图像的亮度
    random_factor = np.random.randint(10, 21) / 10.  # 随机因子
    contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor)  # 调整图像对比度
    random_factor = np.random.randint(0, 31) / 10.  # 随机因子
    return ImageEnhance.Sharpness(contrast_image).enhance(random_factor)  # 调整图像锐度
 
if __name__ == '__main__':
    imageDir = "D:/untitled/数据增强/train/"  # 要改变的图片的路径文件夹
    saveDir = 'D:/untitled/Month/train_aug/' # 要保存的图片的路径文件夹
 
 
    cate = [imageDir + x for x in os.listdir(imageDir)]
    print(cate)
    for idx, folder in enumerate(cate):
        print("idx = ",idx,"folder = ",folder)
        for i in glob.glob(folder + '/*.png'):
            print("i = ",i)
 
            saveName= i[:-4]+"fl.jpg"
            saveImage = flip(i)
            saveImage.save(os.path.join(saveDir,saveName))
 
            saveName = i[:-4] + "ro.jpg"
            saveImage = rotation(i)
            saveImage.save(os.path.join(saveDir, saveName))
 
            saveName = i[:-4] + "be.jpg"
            saveImage = brightnessEnhancement(i)
            saveImage.save(os.path.join(saveDir, saveName))
 
            saveName = i[:-4] + "co.jpg"
            saveImage = contrastEnhancement(i)
            saveImage.save(os.path.join(saveDir, saveName))
 
            saveName = i[:-4] + "cl.jpg"
            saveImage = colorEnhancement(i)
            saveImage.save(os.path.join(saveDir, saveName))
 
            saveName = i[:-4] + "ra.jpg"
            saveImage = randomColor(i)
            saveImage.save(os.path.join(saveDir, saveName))

如果多类别文件夹，如图像分类，需要统一进行数据增强，加入了文件夹遍历，可自取

from tkinter import image_names
from matplotlib.animation import ImageMagickBase


def flip(root_path,img_name):   #翻转图像
    img = Image.open(os.path.join(root_path, img_name))
    filp_img = img.transpose(Image.FLIP_LEFT_RIGHT)
    # filp_img.save(os.path.join(root_path,img_name.split('.')[0] + '_flip.jpg'))
    return filp_img

def rotation(root_path, img_name):
    img = Image.open(os.path.join(root_path, img_name))
    rotation_img = img.rotate(20) #旋转角度
    # rotation_img.save(os.path.join(root_path,img_name.split('.')[0] + '_rotation.jpg'))
    return rotation_img

def randomColor(root_path, img_name): #随机颜色
    """
    对图像进行颜色抖动
    :param image: PIL的图像image
    :return: 有颜色色差的图像image
    """
    image = Image.open(os.path.join(root_path, img_name))
    random_factor = np.random.randint(0, 31) / 10.  # 随机因子
    color_image = ImageEnhance.Color(image).enhance(random_factor)  # 调整图像的饱和度
    random_factor = np.random.randint(10, 21) / 10.  # 随机因子
    brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor)  # 调整图像的亮度
    random_factor = np.random.randint(10, 21) / 10.  # 随机因子
    contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor)  # 调整图像对比度
    random_factor = np.random.randint(0, 31) / 10.  # 随机因子
    return ImageEnhance.Sharpness(contrast_image).enhance(random_factor)  # 调整图像锐度

def contrastEnhancement(root_path, img_name):  # 对比度增强
    image = Image.open(os.path.join(root_path, img_name))
    enh_con = ImageEnhance.Contrast(image)
    contrast = 1.5
    image_contrasted = enh_con.enhance(contrast)
    return image_contrasted

def brightnessEnhancement(root_path,img_name):#亮度增强
    image = Image.open(os.path.join(root_path, img_name))
    enh_bri = ImageEnhance.Brightness(image)
    brightness = 1.5
    image_brightened = enh_bri.enhance(brightness)
    return image_brightened

def colorEnhancement(root_path,img_name):#颜色增强
    image = Image.open(os.path.join(root_path, img_name))
    enh_col = ImageEnhance.Color(image)
    color = 1.5
    image_colored = enh_col.enhance(color)
    return image_colored

from PIL import Image
from PIL import ImageEnhance
import os
import cv2
import numpy as np
imageDir="./image"     #要改变的图片的路径文件夹
saveDir="./save"   #要保存的图片的路径文件夹

path = os.path.join(imageDir)

rootlist=[]
dirlist=[]
filelist=[]
for root, dirnames, filenames in os.walk(path):
    rootlist.append(root)
    for dirname in dirnames:
        dir_save_path = os.path.join(saveDir,dirname)
        if not os.path.exists(dir_save_path):
            os.makedirs(dir_save_path)
        dirlist.append(os.path.join(dir_save_path))
del rootlist[0]

print(dirlist)

for i in range(len(rootlist)):
    for r, d, f in os.walk(rootlist[i]):
        for name in f:
            filelist.append(os.path.join(name))
    filelist.sort(key=lambda x:int(x[:-4]))
    for j in range(len(filelist)):
        saveName= filelist[j][:-4]+"id.png"
        image = Image.open(os.path.join(rootlist[i], filelist[j]))
        image.save(os.path.join(dirlist[i],saveName))

        saveName= filelist[j][:-4]+"fl.png"
        saveImage=flip(rootlist[i], filelist[j])
        saveImage.save(os.path.join(dirlist[i],saveName))

        saveName= filelist[j][:-4]+"ro.png"
        saveImage=rotation(rootlist[i], filelist[j])
        saveImage.save(os.path.join(dirlist[i],saveName))

        saveName= filelist[j][:-4]+"be.png"
        saveImage=brightnessEnhancement(rootlist[i], filelist[j])
        saveImage.save(os.path.join(dirlist[i],saveName))

        saveName= filelist[j][:-4]+"ra.png"
        saveImage=randomColor(rootlist[i], filelist[j])
        saveImage.save(os.path.join(dirlist[i],saveName))

        saveName= filelist[j][:-4]+"con.png"
        saveImage=contrastEnhancement(rootlist[i], filelist[j])
        saveImage.save(os.path.join(dirlist[i],saveName))    

        saveName= filelist[j][:-4]+"col.png"
        saveImage=colorEnhancement(rootlist[i], filelist[j])
        saveImage.save(os.path.join(dirlist[i],saveName))
    
    filelist.clear()