使用经过训练的 TensorFlow model 预测单个 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)
]) 

This is the code for the model, which I have trained using the mnist dataset. What I want to do is to then pass a 28x28 png image to the predict() method, which is not working. The code for the prediction is:

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

which produces the error

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

I have been stuck on this problem for a few days, but I can't find the correct way to pass an image into the predict method. Any help?

Answer 1

Predict function makes predictions over a batch of image. You should include batch dimension (first dimension) to your img, even to predict a single example. You need something like this:

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

As @desertnaut says, seems you are using a RGB image, so your first layer should use input_shape = (28,28,3) . Therefore, img parameter of predict function should have (1,28,28,3) shape.

In your case, img parameter of predict function has (28,28,3) shape, thus predict function took the first dimension as number of images, and could not match the other two dimensions to the input_shape of the first layer.