Pytorch中利用Pillow随机增强数据集

最近在收集数据图片作为数据集时常常会遇到一个问题，就是图片数据的数量不够，这样会导致神经网络的训练效果不够理想，因此在查阅了相关资料之后，发现python中的pillow函数可以实现对图片的各种处理，如翻转、旋转、改变亮度和锐度等等，但是仅仅做其中某个操作往往显得较为单一，我更希望的是能够将这些处理图片的操作综合起来，并且进行随机综合处理，故写下如下代码便于自用，现将其分享，方便大家快速处理数据集，也欢迎讨论。

"""
Author:Keycle
Date:2020/11/07
随机扩充数据集
注意图片格式只能是jpg
"""

# -*- coding: utf-8 -*-

import os
import time
import random
from PIL import Image,ImageChops,ImageEnhance

def image_enhanced(img,savefilepath,save_filename):
    """图像随机扩充"""
    rand_choice = random.randrange(0, 2)   # 随机左右翻转
    rand = random.randrange(-60, 60, 10)    # 选择旋转角度-60~60（逆时针旋转）
    randbri = random.choice([0.6, 0.8, 1.0, 1.2])   # 选择亮度，大于1增强，小于1减弱
    randcol = random.choice([0.7, 0.9, 1.1, 1.3])   # 选择色度，大于1增强，小于1减弱
    randcon = random.choice([0.7, 0.9, 1.1, 1.3])   # 选择对比度，大于1增强，小于1减弱
    randsha = random.choice([0.5, 1.0, 1.5])        # 选择锐度，大于1增强，小于1减弱
    if rand_choice == 0:
        lr = img.transpose(Image.FLIP_LEFT_RIGHT)
        out1=lr.rotate(rand)
        bri = ImageEnhance.Brightness(out1)
        bri_img1 = bri.enhance(randbri)
        col = ImageEnhance.Color(bri_img1)
        col_img1 = col.enhance(randcol)
        con = ImageEnhance.Contrast(col_img1)
        con_img1 = con.enhance(randcon)
        sha = ImageEnhance.Sharpness(con_img1)
        sha_img1 = sha.enhance(randsha)
        sha_img1.save(savefilepath + save_filename)
    elif rand_choice == 1:
        out1 = img.rotate(rand)
        bri = ImageEnhance.Brightness(out1)
        bri_img1 = bri.enhance(randbri)
        col = ImageEnhance.Color(bri_img1)
        col_img1 = col.enhance(randcol)
        con = ImageEnhance.Contrast(col_img1)
        con_img1 = con.enhance(randcon)
        sha = ImageEnhance.Sharpness(con_img1)
        sha_img1 = sha.enhance(randsha)
        sha_img1.save(savefilepath + save_filename)


# 定义扩充图片函数
def image_expansion(filepath,savefilepath,save_prefix):
    """
    :param filepath: 图片路径
    :param savefilepath: 扩充保存图片路径
    :param save_prefix: 图片前缀
    :return: 图片扩充数据集
    """
    i = 1
    for parent, dirnames, filenames in os.walk(filepath):
        for filename in filenames:
            image_path=filepath+filename
            print('正在扩充图片：%s' %filename)
            try:
                cycle=1
                img=Image.open(image_path)
                if img.mode == "P":
                    img = img.convert('RGB')

                '''cycle数即为单张图片扩充的图片数'''
                while(cycle<=7):
                    image_enhanced(img,savefilepath,save_filename=save_prefix+str(i)+'.jpg')
                    i += 1
                    cycle += 1

            except Exception as e:
                print(e)
                pass



if __name__ == '__main__':
    # 设置图片路径
    filepath = 'F:/Machine Learning/dataset/test/'  # 注意最后不要少了斜杠

    # 设置扩充保存图片路径
    savefilepath ='F:/Machine Learning/testcases/trainB/'  # 注意最后不要少了斜杠

    # 设置前缀图片名称
    save_prefix='enhanced_0_'

    time1 = time.time()
    image_expansion(filepath, savefilepath,save_prefix)
    time2 = time.time()
    print('总共耗时：' + str(time2 - time1) + 's')

代码运行效果如下：

下面是随机增强后的图片：





可以看到对原图片进行随机进行了翻转、旋转和调整亮度、锐度等操作，当对图片批量操作后，即可实现随机增强数据集