大图像数据集下，如何在循环中将predict_generator输出写入文件？

我来帮你搞定这个问题！核心思路就是分批次处理图像，每处理完一批就把预测结果追加写入同一个文件，这样既不会占用太多内存，最后也能得到完整的结果集。下面给你具体的实现步骤和代码示例：

关键注意点先明确 #

首先你要保证生成器的shuffle=False，这个非常重要！它能确保生成器输出的图像顺序和filenames列表的顺序完全对应，避免预测结果和文件名匹配错乱。

完整实现代码 #

import csv
import numpy as np
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Flatten, Dense, Dropout

# 替换成你的实际参数
img_width, img_height = 224, 224  # 根据你的模型输入尺寸调整
path = "你的验证集目录路径"
top_model_weights_path = "你的模型权重文件路径"

# 构建图像生成器（务必保持shuffle=False）
validation_generator = ImageDataGenerator(rescale=1./255).flow_from_directory(
    path,
    target_size=(img_width, img_height),
    batch_size=6,
    shuffle=False
)
print("生成器构建完成")

# 加载你的模型（注意顺序：先定义模型再加载权重，再做预测）
model = Sequential()
model.add(Flatten(input_shape=(3, 3, 1536)))
model.add(Dense(256, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(6, activation='softmax'))
model.load_weights(top_model_weights_path)
print("顶层模型权重加载完成")

# 初始化结果文件：先写入表头
output_file = "predictions.csv"
with open(output_file, 'w', newline='') as f:
    writer = csv.writer(f)
    # 表头包含：文件名、每个类别的概率、预测类别
    header = ["filename"] + [f"class_{i}_proba" for i in range(6)] + ["predicted_class"]
    writer.writerow(header)

# 计算总处理步数：总样本数除以批次大小，向上取整
total_steps = np.ceil(validation_generator.samples / validation_generator.batch_size).astype(int)

# 循环处理每一批次
for step in range(total_steps):
    # 获取当前批次的图像（生成器next()返回图像和标签，我们只需要图像）
    batch_images, _ = next(validation_generator)
    # 生成当前批次的预测结果
    batch_proba = model.predict_proba(batch_images)
    batch_classes = model.predict_classes(batch_images)
    
    # 获取当前批次对应的文件名：从生成器的filenames中切片，处理最后一批可能不足批次大小的情况
    start_idx = step * validation_generator.batch_size
    end_idx = min(start_idx + validation_generator.batch_size, validation_generator.samples)
    batch_filenames = validation_generator.filenames[start_idx:end_idx]
    
    # 将当前批次结果追加写入CSV文件
    with open(output_file, 'a', newline='') as f:
        writer = csv.writer(f)
        for filename, proba, cls in zip(batch_filenames, batch_proba, batch_classes):
            # 把概率数组转成列表，和文件名、类别一起组成一行数据
            row = [filename] + list(proba) + [cls]
            writer.writerow(row)
    
    print(f"已完成第 {step+1}/{total_steps} 批处理")

print(f"所有预测结果已保存到 {output_file}")

后续读取文件计算概率的方法 #

保存成CSV格式后，后续读取和计算会非常方便，比如用pandas快速统计：

import pandas as pd

# 读取结果文件
df = pd.read_csv("predictions.csv")

# 示例1：计算每个类别的平均概率
class_avg_proba = df[[f"class_{i}_proba" for i in range(6)]].mean()
print("每个类别的平均概率：")
print(class_avg_proba)

# 示例2：统计每个类别的预测数量
pred_class_counts = df["predicted_class"].value_counts()
print("\n每个类别的预测样本数：")
print(pred_class_counts)

额外提醒 #

1. 你原来的代码里有个顺序错误：先调用了predict_generator但模型还没加载，上面的代码已经修正了这个问题，确保先加载模型再做预测。
2. 如果你的Keras版本比较新，predict_proba和predict_classes可能已经被弃用，可以改用predict方法，然后自己提取概率和类别：

   batch_preds = model.predict(batch_images)
   batch_proba = batch_preds  # 直接就是概率数组
   batch_classes = np.argmax(batch_preds, axis=1)  # 取概率最大的索引作为类别
   

内容的提问来源于stack exchange，提问作者Rehan Aziz