keras文件读取

import tensorflow as tf
import tensorflow.keras as keras
import matplotlib.pyplot as plt
import os
import PIL
import pathlib
import math
import random
import numpy as np
import shutil
import PIL
# 划出测试图像
def div_train_test(data_dir):
    data_dir = pathlib.Path(data_dir)
    image_count = len(list(data_dir.glob('*/*')))  # 读取文件夹下面所有的文件
    print('一共有有{}张图像'.format(image_count))
    # -----------------------------------------------------------------------
    # 查看文件有多少类，并输出第一个类文件中的第一张图像
    # 且查看第一张图像的size和通道数
    name_list = []
    for item in data_dir.iterdir():
        name_list.append(item.name)
        # print(item.resolve())
        # print(item)

    print('有如下类别{}共{}类'.format(name_list, len(name_list)))
    # --------------------------------------------------------------------------
    # 选出百分之30的图像放入测试文件夹中
    n = len(name_list)
    path_join = ('test')
    if os.path.exists(path_join):
        pass
    else:
        os.mkdir(path_join)

    for i in range(n):
        tar_path = os.path.join(path_join, name_list[i])
        if os.path.exists(tar_path):
            pass
        else:
            os.mkdir(tar_path)
        so_path = os.path.join(data_dir, name_list[i])
        so_path = pathlib.Path(so_path)
        num = len(list(so_path.glob('*/')))
        print(num)
        test_num = math.ceil(num * 0.3)
        rd = random.sample(range(0, num), test_num)
        print(len(rd))
        name = []
        path_join = pathlib.Path(path_join)
        for item in list(so_path.glob('*/')):
            name.append(item.name)
        if so_path.is_dir():
            print('正在从{}拷贝'.format(so_path))
            for j in (rd):
                im = PIL.Image.open(str(so_path) + '\\' + name[j])
                im.save(tar_path + '\\' + name[j])
                os.remove(str(so_path) + '\\' + name[j])  # 转移完后删除原图片


# --------------------------------------------------------------------------




def datapreprocess(data_dir,batch_size,img_height,img_width,data_augment):
    '''

    Args:
        data_dir:文件路径
        batch_size: batch大小
        img_height: 图像高度
        img_width: 图像宽度
        data_augment:是否数据增强,True或False,默认False
    Returns:tain_ds和val_ds分别训练和验证集

    '''
    #读取文件，查看相关的属性
    # --------------------------------------------
    # 如果test下面文件为空则划分出测试集
    flag=os.getcwd()
    flag=os.path.join(flag,'test')
    flag=pathlib.Path(flag)
    if len(list(flag.glob('*/*')))==0:
        div_train_test('train')
    # --------------------------------------------


    print(flag)
    data_dir=pathlib.Path(data_dir)
    image_count=len(list(data_dir.glob('*/*'))) #读取文件夹下面所有的文件
    print('一共有有{}张图像'.format(image_count))
# -----------------------------------------------------------------------
    # 查看文件有多少类，并输出第一个类文件中的第一张图像
    # 且查看第一张图像的size和通道数
    name_list=[]
    for item in data_dir.iterdir():
        name_list.append(item.name)
        # print(item.resolve())
        # print(item)

    print('有如下类别{}共{}类'.format(name_list,len(name_list)))

    ph1_dir=os.path.join(data_dir,name_list[0])
    ph1_dir=pathlib.Path(ph1_dir)
    ph1=list(ph1_dir.glob('*/'))
    p=PIL.Image.open(ph1[0])
    plt.imshow(p)
    plt.colorbar()
    plt.show()
    #查看图像的shape
    p=tf.io.read_file(str(list(data_dir.glob('*/*'))[0]))
    p_tensor=tf.image.decode_image(p)
    print(p_tensor.shape)

# -----------------------------------------------------------------------
    #利用keras.preprocessing来创建数据集
    # 因为默认的color_mode参数是‘rgb’，所以如果是灰度图则需要将color_mode改为‘grayscle’
    if str(data_dir)=='train':
        train_ds = keras.preprocessing.image_dataset_from_directory(
            data_dir, validation_split=0.2, subset='training', seed=111,
            image_size=[img_height, img_width], color_mode='grayscale', batch_size=batch_size)
        print(train_ds)
        val_ds = tf.keras.preprocessing.image_dataset_from_directory(
            data_dir, validation_split=0.2, subset='validation', seed=111,
            image_size=[img_height, img_width], color_mode='grayscale', batch_size=batch_size
        )
        print(val_ds)
    if str(data_dir) == 'test':
        test_ds = keras.preprocessing.image_dataset_from_directory(
            data_dir, image_size=[img_height, img_width], color_mode='grayscale')

    # --------------------------------------------------
    #显示图像的类名和显示图像
    # class_names=train_ds.class_names
    class_names=name_list
    print(class_names) #此方法也可以输出类别名
    plt.figure(figsize=(10, 10))
    if str(data_dir)=='train':
        for images, labels in train_ds.take(1):
            for i in range(9):
                ax = plt.subplot(3, 3, i + 1)
                plt.imshow(images[i].numpy().astype("uint8"))
                plt.title(class_names[labels[i]])
                plt.axis("off")
    # 查看图像的形状
    if str(data_dir)=='train':
        for image_batch, labels_batch in train_ds:
            print(image_batch.shape)
            print(labels_batch.shape)
            break
    # --------------------------------------------------
    #利用prefetch和cache加速数据的读取
    # ---------------------------------------------------
    AUTOTUNE=tf.data.experimental.AUTOTUNE
    if str(data_dir)=='train':
        train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
        val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
    if str(data_dir)=='test':
        test_ds=test_ds.cache().prefetch(buffer_size=AUTOTUNE)
    #标准化数据
    normalization_layer = keras.layers.experimental.preprocessing.Rescaling(1. / 255)
    if str(data_dir)=='train':
        train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y))
        val_ds=val_ds.map(lambda x, y: (normalization_layer(x), y))
    if str(data_dir)=='test':
        test_ds=test_ds.map(lambda x, y: (normalization_layer(x), y))
    # image_batch, labels_batch = next(iter(normalized_ds))
    # first_image = image_batch[0]
    # Notice the pixels values are now in `[0,1]`.
    # print(np.min(first_image), np.max(first_image))

        # plt.figure(figsize=(10,10)) #显示被增强的图像
        # for images,_ in train_ds.take(1):
        #     for i in range(9):
        #         augmented_image=data_augmentation(images)
        #         ax=plt.plot(3,3,i+1)
        #         plt.imshow(augmented_image[0].numpy().astype('uint8'))
        #         plt.show()
    if data_augment:
        # 数据增强处理
        data_augmentation = keras.Sequential([
            # 因为numpy版本问题出错，flip不用
            # keras.layers.experimental.preprocessing.RandomFlip('horizontal',input_shape=[img_height, img_width, p_tensor.shape[2]]),
            keras.layers.experimental.preprocessing.RandomRotation(0.1),
            keras.layers.experimental.preprocessing.RandomZoom(0.1),
        ])
    if str(data_dir)=='train':
        print('当前data_dir is ',data_dir)
        if data_augment:
            return train_ds, val_ds, data_augmentation
        # else:
        #     return train_ds, val_ds
    if str(data_dir) == 'test':
        print('当前data_dir is ', data_dir)
        return test_ds
    # if str(data_dir)=='train':
    #     re=(train_ds,val_ds,data_augmentation)
    #     print('re is ',re)
    #     return re
    # if str(data_dir)=='test':
    #     print('test is ',test_ds)
    #     return test_ds




train,val,augmentation=datapreprocess('train',32,224,224,True)
test=datapreprocess('test',32,224,224,False)
print(train,val)
print(test)