利用keras-ImageDataGenerator实现图像增强与批量生成

---

1.背景

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

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

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()
        """
        get_pic_path返回list,该list为2维，存储格式为:(类，该类下图片路径)。
        例如：list[0]:表示第一类下所有图片的路径;list[0][0]:表示第一类下，第一张图片的绝对路径.
        """
        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 self.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')#data_increase为存放分好类的待增强的图片数据的文件夹
    p = picture_increating(path)
    p.increate_pic()

2.代码解释

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

3.演示

3.1 具体操作

1. 初始,准备数据,创建空.py文件
   

2. 将代码粘贴到un.py文件中(可任意命名),且该文件需要与example同级
   

3. 修改path中data_increase为example
   

4. 打开cmd,cd到该文件夹下.(此处为C:\Users\6\Desktop\新建文件夹)
   

5. 在控制台内输入python un.py即完成
   
   

3.2 补充

 初始图片有8类,每类1~6张

带入路径:C:\Users\6\Desktop\taidi_cup\code\data_increase,在os.getcwd()路径下生成文件夹dataset并给每类创建新文件并夹在新文件夹中生成图片.

图片个数=20*原类所含图片个数

注意：工作路径要改为数据文件夹的上一级，此文章是在工作路径：C:\Users\6\Desktop\taidi_cup\code下进行的。

4.总结

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