使用flow_from_directory时,如何从Keras模型预测的独热编码向量确定类别?
如何从Keras独热编码向量中匹配实际类别
你遇到的这个问题很典型——用flow_from_directory自动生成类别向量后,确实需要一个方法把模型输出的独热编码和实际的类别名称对应起来,下面给你一步步讲清楚怎么解决:
1. 先拿到类别与索引的映射关系
flow_from_directory返回的生成器对象自带一个class_indices属性,它本质是个字典,键是你的文件夹名称(也就是类别名),值是对应的类别索引。比如你的文件夹结构里,它会生成类似 {'cow': 0, 'pig': 1} 的映射(默认按文件夹名称的字母顺序排序)。
你可以直接打印这个属性来确认:
print(train_generator.class_indices)
2. 把独热编码向量转成类别名称
当模型输出独热编码向量(比如 [0.97, 0.03]),我们需要先找到向量中最大值对应的索引,再通过上面的映射字典反向找到类别名称。这里有两种简单的实现方式:
方式一:反转映射字典(推荐)
先把class_indices反转成“索引→类别名”的字典,之后就能快速查找:
# 反转class_indices,得到索引到类别的映射 idx_to_class = {v: k for k, v in train_generator.class_indices.items()}
然后用np.argmax()获取独热向量的最大值索引,再通过反转后的字典拿到类别:
import numpy as np # 假设pred是模型预测得到的独热编码,形状为(1, 类别数) pred = model.predict(your_input_image) predicted_idx = np.argmax(pred, axis=1)[0] predicted_class = idx_to_class[predicted_idx] print(f"预测的类别是:{predicted_class}")
方式二:直接遍历原字典查找
如果不想额外创建反转字典,也可以直接遍历class_indices找到对应类别:
predicted_idx = np.argmax(pred, axis=1)[0] predicted_class = [k for k, v in train_generator.class_indices.items() if v == predicted_idx][0]
3. 额外注意点
- 要保证训练和测试生成器的类别索引一致:只要你的
train和test文件夹下的类别子文件夹名称完全相同,flow_from_directory生成的索引就会一致。如果担心排序问题,可以手动指定classes参数固定类别顺序:
train_generator = train_datagen.flow_from_directory( '/train/', target_size=(x, y), batch_size=batch_size, class_mode='categorical', classes=['cow', 'pig'] # 手动指定类别顺序,确保训练、测试生成器的索引完全匹配 )
这样就能轻松把模型输出的独热编码和实际的类别名称对应起来啦!
内容的提问来源于stack exchange,提问作者Joe
相关产品推荐
相关产品推荐

