模仿 ImageDataGenerator().flow_from_directory() 实现的图片数据集导入，但返回的是数据与标签

import os
import numpy as np
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.preprocessing.image import img_to_array, load_img


# 多类独热码
def load_dataset_CNN(base_dir, target_size, color_mode='rgb', shuffle=True):
    X, y = [], []
    total_num = 0

    classes = os.listdir(base_dir)
    for clas in classes:
        label = classes.index(clas)
        imgs_dir = os.path.join(base_dir, clas)
        imgs_names = os.listdir(imgs_dir)

        for img_name in imgs_names:
            img_path = os.path.join(imgs_dir, img_name)
            img_PIL = load_img(path=img_path, target_size=target_size, color_mode=color_mode)
            img_array = img_to_array(img_PIL)
            X.append(img_array)
            y.append(label)

        print("Found {} images in class '{}', setting as class[{}].".format(len(imgs_names), clas, label))
        total_num += len(imgs_names)
    print('{} images in total.'.format(total_num))

    X = np.array(X)
    y = to_categorical(np.array(y))

    if shuffle is True:
        state = np.random.get_state()
        np.random.shuffle(X)
        np.random.set_state(state)
        np.random.shuffle(y)

    return X, y

🌰：

from load_dataset import load_dataset_CNN

IMG_HEIGHT, IMG_WIDTH = 100, 100
X, y = load_dataset_CNN(base_dir='./train', target_size=(IMG_HEIGHT, IMG_WIDTH))
print(X)
print(y)

…
👍