以cifar10为例:
import os
import tensorflow as tf
import matplotlib.pyplot as plt
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.cifar10.load_data()
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
index_to_name = {i: name for i, name in enumerate(class_names)}
#根据class_name创建不同的文件夹
for index,name in index_to_name.items():
path = os.path.join('./data/train', name)
os.makedirs(path, exist_ok=True)
for index,name in index_to_name.items():
path = os.path.join('./data/test', name)
os.makedirs(path, exist_ok=True)
#读取cifar10数据集,并将其保存到对应的文件夹中
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.cifar10.load_data()
for i in range(train_images.shape[0]):
img = train_images[i]
label = train_labels[i][0]
name = index_to_name[label]
path = os.path.join('./data/train', name, f'{i}.png')
plt.imsave(path, img)
for i in range(test_images.shape[0]):
img = test_images[i]
label = test_labels[i][0]
name = index_to_name[label]
path = os.path.join('./data/test', name, f'{i}.png')
plt.imsave(path, img)