数据增强实现

为了训练小样本的模型，通常需要对训练集进行数据增强，下面是数据增强的实现代码：

from PIL import Image
from PIL import ImageEnhance
import os
import numpy as np

def changeImg(root_path,img_name,angle,contrast,brightness,color):

    img = Image.open(os.path.join(root_path, img_name))
    filp_img = img.transpose(Image.FLIP_LEFT_RIGHT)  #翻转图像


    rotation_img = img.rotate(angle)  # 旋转角度
    filp_rot_img = filp_img.rotate(angle)

    random_factor = np.random.randint(0, 31) / 10.  # 随机因子
    color_image = ImageEnhance.Color(img).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.  # 随机因子
    sharp_img = ImageEnhance.Sharpness(contrast_image).enhance(random_factor)  # 调整图像锐度


    enh_con = ImageEnhance.Contrast(img)
    image_contrasted = enh_con.enhance(contrast)  #对比度增强


    enh_bri = ImageEnhance.Brightness(img)
    image_brightened = enh_bri.enhance(brightness) #亮度增强


    enh_col = ImageEnhance.Color(img)
    image_colored = enh_col.enhance(color)  #颜色增强

    return filp_img,rotation_img,filp_rot_img,sharp_img,image_contrasted,image_brightened,image_colored,color_image,brightness_image,contrast_image

imageDir="./old/" #要改变的图片的路径文件夹
saveDir="./new/"   #要保存的图片的路径文件夹
for j in range(5):
    for name in os.listdir(imageDir):
        path = name.split('.')[0]
        angle = np.random.randint(0,20)
        contrast = np.random.uniform(1,2)
        brightness = np.random.uniform(1,2)
        color = np.random.uniform(1,2)
        images = changeImg(imageDir,name,angle,contrast,brightness,color)
        for i in range(len(images)):
            saveName = path + '_' + str(i)+ '_' + str(j) + ".jpg"
            images[i].save(os.path.join(saveDir, saveName))

只需要更改读取图片的路径imageDir和保存新的图片的路径 saveDir，for j in range(5)表示对同一张图片进行5次增强。