import requests
from io import BytesIO
from matplotlib import pyplot as plt

image_url = "https://datasets-server.huggingface.co/assets/mnist/--/mnist/train/12/image/image.jpg"

# Download the image using requests
response = requests.get(image_url)

if response.status_code == 200:
     # Read the image from the response content
    image_bytes = BytesIO(response.content)
    input_image = plt.imread(image_bytes, format='jpg')

    # Display the downloaded image
    plt.imshow(input_image, cmap='gray')
    plt.show()

    # Preprocess the image (assuming the model expects the same preprocessing as the MNIST dataset)
    input_image = np.array(input_image, dtype='float32')
    input_image /= 255.0  # Normalize pixel values to the range [0, 1]

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

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

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

    print("Predicted Class:", predicted_class)
else:
    print("Failed to download the image. Status code:", response.status_code)