from matplotlib import pyplot as plt
import numpy as np

# Assuming x_test has shape (num_samples, 28, 28)
index_to_predict = 13
image_to_predict = x_test[index_to_predict]

input_image = np.array(x_test[index_to_predict], dtype='float')
pixels = input_image.reshape((28, 28))
plt.imshow(pixels, cmap='gray')
plt.show()

print("Seeing if we can predict that it's " + str(y_test[index_to_predict]))

# Reshape the image to match the model's input shape
image_to_predict = tf.reshape(image_to_predict, (1, 28, 28))

# Get the prediction using the probability_model
predictions = probability_model.predict(image_to_predict)

# The predictions will be a probability distribution over classes.
# You can extract the class with the highest probability using argmax.
predicted_class = tf.argmax(predictions[0]).numpy()

print("Predicted Class:", predicted_class)