我有两个图片文件夹，一个文件夹包含8296张窃电用户图像，另外一个文件夹包含8296张非窃电图像。请帮我写出用ImageDataGenerator生成数据集的代码，并且按照8：1：1比例划分训练集、验证集和测试集

好的，以下是代码实现：

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`方法从文件夹中读取图像数据，并进行数据增强和预处理。接着，我们根据给定的比例划分了训练集、验证集和测试集，并将它们保存到对应的文件夹中，同时将文件名和标签保存到文本文件中，以备后续使用。