python - 使用经过训练的 TensorFlow 模型预测单个 PNG 图像

Question

import tensorflow as tf
model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape = (28,28)),
    tf.keras.layers.Dense(128, activation = 'relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
]) 

这是模型的代码，我使用 mnist 数据集对其进行了训练。我想做的是将 28x28 png 图像传递给 Predict() 方法，该方法不起作用。预测代码为:

img = imageio.imread('image_0.png')
prediction = model.predict(img, batch_size = 1)

产生错误

ValueError: Error when checking input: expected flatten_input to have shape (28, 28) but got array with shape (28, 3)

我已经被这个问题困扰了几天，但我找不到将图像传递到预测方法的正确方法。有什么帮助吗？

Answer 1

预测函数对一批图像进行预测。您应该在您的 img 中包含批处理维度(第一维度)，甚至是为了预测单个示例。 你需要这样的东西:

img = imageio.imread('image_0.png')
img = np.expand_dims(img, axis=0)
prediction = model.predict(img)

正如@desertnaut所说，似乎您使用的是RGB图像，因此您的第一层应该使用input_shape = (28,28,3)。因此，预测函数的 img 参数应具有 (1,28,28,3) 形状。

在您的情况下，预测函数的 img 参数具有 (28,28,3) 形状，因此预测函数将第一个维度作为图像数量，并且无法将其他两个维度与第一层的input_shape。