背景

最近在做一个图像题,初始样本图给的特别少(每类只有几张图),因此想通过keras-ImageDataGenerator去做数据增强,但是这样的数据带入模型训练泛化性也是很差的,只是先做着,等数据全部公布后再做.在此之前,先将数据增强函数写好,到时候方便使用,以下为具体代码:

import torch
import pandas as pd
import cv2
import os
from PIL import Image
import matplotlib.pyplot as plt
import pandas as pd
import copy
import numpy as np
from keras.preprocessing import image
from keras.preprocessing.image  import ImageDataGenerator 

class picture_increating:
    """
    利用keras中的ImageDataGenerator做图像增强;
    代入各类所在文件夹的父文件夹路径,即可在与该父文件同级路径下创建各类文件夹,并生成经过处理(旋转\剪切\缩放等)的图片.
    """
    def __init__(self,path):
#         self.data_path = os.path.join(os.getcwd(),'data_increase')
        self.data_path = path
        self.datagen =ImageDataGenerator(
            rotation_range = 40,
            width_shift_range=0.2,

            height_shift_range=0.2,
            shear_range =0.2,
            zoom_range=0.2,
        #     horizontal_flip=True,
            fill_mode='reflect'#nearest;reflect;constant;wrap
        )
    #取各图像路径
    def get_pic_path(self):
        """
        代入各类文件夹的父文件夹路径,返回各子类图片路径构成的list
        """
        pic_path=[]
        pic_d =[]
        for i in os.listdir(self.data_path):
            for j in os.listdir(os.path.join(self.data_path,i)):
                pic_path.append(os.path.join(self.data_path,i,j))
            pic_d.append(pic_path)
            pic_path=[]
        return pic_d
    
    def new_package(self,path):
        isExists=os.path.exists(path)
        if not isExists:
            # 如果不存在则创建目录
             # 创建目录操作函数
            os.makedirs(path) 
            print(path+' 创建成功')
        else:
            # 如果目录存在则不创建，并提示目录已存在
            print (path+' 目录已存在')
    
    def increate_pic(self):
        ppath = self.get_pic_path()
        for i in ppath:
            for j in i:
                img = image.load_img(j,target_size=(150,150))
        #         print(type(img))
                x = image.img_to_array(img)
        #         print(x.shape)
                x = x.reshape((1,)+x.shape)#数据扩维度

                ii= 0
                for batch in datagen.flow(x,batch_size=2):
                    #保存的文件名编写
                    lei =j.split('\\')[-2]
                    pic_name = j.split('\\')[-1].split('_')[-1]
                    pic_num = j.split('\\')[-1].split('_')[1]
                    name=lei+pic_num+'_'+str(ii)+'_'+pic_name
                    # 创建新文件夹
                    self.new_package(os.path.join(os.getcwd(),'dataset',lei))
                    
                    
                    #写入图片保存
                    pict = image.array_to_img(batch[0])
#                     pict.save(os.path.join(data_path,lei,name))
                    pict.save(os.path.join(os.getcwd(),'dataset',lei,name))

    
#                     展示
#                     plt.figure(i)
#                     imgplot=plt.imshow(pict)
#                     print(batch.shape,batch[0].shape)
#                     print(type(batch),type(image.array_to_img(batch[0])))
#                     cv2.imwrite(str(i)+'c.jpg',batch[0])

                    # 迭代停止条件
                    ii+=1
                    if ii %20==0:
                        break
                        
if __name__ =='__main__':
    path=os.path.join(os.getcwd(),'data_increase')
    p = picture_increating(path)
    p.increate_pic()

代码解释

代码目的:使用者只代入各类所在文件夹的父文件夹路径,即可在与该父文件同级路径下创建各类文件夹,并生成经过处理(旋转\剪切\缩放等)的图片.

初始定义self.data_path和ImageDataGenerator类对象self.datagen
get_pic_path:用于取各类图片的路径,并存为list返回
new_package:用于创建父文件同级的新文件dataset
increate_pic(核心):即为数据迭代生成函数.上面的函数都在该函数内使用.要理解该函数核心是理解ImageDataGenerator.flow(),之后可通过更改参数变换图像,此处不赘诉.

图片演示

初始图片有8类,每类1~6张
在这里插入图片描述

带入路径:C:\Users\6\Desktop\taidi_cup\code\data_increase
生成文件夹dataset并给每类创建新文件夹在新文件夹中生成图片.
图片个数为20*原类所含图片个数

总结:

ImageDataGenerator对象不仅有.flow还有.flow_from_directory和.flow_DataFrame等,区别在于带入路径参数与返回值不一样,.flow_from_directory也能实现批量生成,但是生成路径不好自定义,因此便自己利用.flow写了该类.欢迎讨论。