我有两个图片文件夹,一个文件夹包含8296张窃电用户图像,另外一个文件夹包含8296张非窃电图像。请帮我写出用ImageDataGenerator生成数据集的代码,并且按照8:1:1比例划分训练集、验证集和测试集
时间: 2024-03-03 13:50:11 浏览: 160
将图片和标注数据按比例切分为训练集和测试集
好的,以下是代码实现:
```python
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import os
# 定义路径和参数
train_dir = 'train'
val_dir = 'val'
test_dir = 'test'
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1
# 创建文件夹
if not os.path.exists(train_dir):
os.makedirs(train_dir)
if not os.path.exists(val_dir):
os.makedirs(val_dir)
if not os.path.exists(test_dir):
os.makedirs(test_dir)
# 图像数据生成器
datagen = ImageDataGenerator(rescale=1./255, validation_split=val_ratio+test_ratio)
# 训练集生成器
train_generator = datagen.flow_from_directory(
'dataset',
target_size=(224, 224),
batch_size=32,
class_mode='binary',
subset='training',
save_to_dir=train_dir)
# 验证集生成器
val_generator = datagen.flow_from_directory(
'dataset',
target_size=(224, 224),
batch_size=32,
class_mode='binary',
subset='validation',
save_to_dir=val_dir)
# 测试集生成器
test_generator = datagen.flow_from_directory(
'dataset',
target_size=(224, 224),
batch_size=32,
class_mode='binary',
subset='validation',
save_to_dir=test_dir)
# 计算样本数量
total_samples = train_generator.samples
train_samples = int(total_samples * train_ratio)
val_samples = int(total_samples * val_ratio)
test_samples = int(total_samples * test_ratio)
# 保存数据集
train_images, train_labels = train_generator.next()
val_images, val_labels = val_generator.next()
test_images, test_labels = test_generator.next()
# 将数据写入文件
with open('train.txt', 'w') as f:
for i in range(train_samples):
f.write(train_dir + '/' + str(i+1) + '.jpg' + ' ' + str(int(train_labels[i])) + '\n')
with open('val.txt', 'w') as f:
for i in range(val_samples):
f.write(val_dir + '/' + str(i+1) + '.jpg' + ' ' + str(int(val_labels[i])) + '\n')
with open('test.txt', 'w') as f:
for i in range(test_samples):
f.write(test_dir + '/' + str(i+1) + '.jpg' + ' ' + str(int(test_labels[i])) + '\n')
```
代码中,我们使用了`ImageDataGenerator`生成器,通过`flow_from_directory`方法从文件夹中读取图像数据,并进行数据增强和预处理。接着,我们根据给定的比例划分了训练集、验证集和测试集,并将它们保存到对应的文件夹中,同时将文件名和标签保存到文本文件中,以备后续使用。
阅读全文