如何基于CNN AI预测结果展示对应图片？

错误原因拆解 #

• 第一个错误（plt.imshow(test_set[location])）：你的test_set是批量数据集（比如Keras ImageDataGenerator生成的对象或TensorFlow Dataset），直接通过索引取到的是整个批次的图片数组（形状(32,224,224,3)），而plt.imshow需要单张图片的形状（(224,224,3)），所以报错。
• 第二个错误（plt.imshow(classes[location])）：classes是模型输出的预测概率数组（形状(样本数, 类别数)，比如二分类时是(N,2)），这不是图片数据，自然无法用imshow显示。

修正代码 #

场景1：test_set是numpy数组（形状为(样本总数, 224,224,3)） #

直接提取单张图片即可：

location = 2
# 打印预测类别
print(categories[np.argmax(classes[location])])
# 获取单张图片
img = test_set[location]
# 若预处理时做了归一化（比如除以255），需还原为0-255的整数
# img = img * 255
# img = img.astype(np.uint8)
plt.figure()
plt.imshow(img)
plt.show()

场景2：test_set是批量数据集（如ImageDataGenerator、TensorFlow Dataset） #

需要先从批次中提取图片数组，再取单张：

location = 2
# 获取测试集的一个批次（若location超过批次大小，需循环获取对应批次）
test_batch_images, _ = next(iter(test_set))
# 提取单张图片
img = test_batch_images[location]
# 打印预测类别
print(categories[np.argmax(classes[location])])
plt.figure()
plt.imshow(img)
plt.show()

关键注意事项 #

• 若测试图片在预处理阶段做了归一化操作，显示前必须还原到0-255的整数范围，否则plt.imshow会显示异常。
• classes仅用于判断预测类别，图片必须从原始测试集中获取，二者通过相同的索引location对应匹配。

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