import gradio as gr
import numpy as np
from PIL import Image  # Import the Image module from PIL

# Assuming your model is loaded as 'loaded_model'

def predict_image(image):
    """Predicts the class of an image using the loaded model."""
    # Resize the image to the desired shape (28, 28)
    image = Image.fromarray(image).resize((28, 28))  
    image = np.array(image)  # Convert back to NumPy array
    # Preprocess the image (normalize, etc.) if necessary
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    prediction = loaded_model.predict(image)
    # Process the prediction (e.g., get the class with highest probability)
    predicted_class = np.argmax(prediction)
    return predicted_class

# Create a Gradio interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(),  # Input image (no shape specified)
    outputs="label",  # Output label
    title="MNIST Digit Classifier",
    description="Upload an image of a handwritten digit (0-9) to classify it."
)

iface.launch()