#denis_mnist_cnn_model_resnet50_v1.h5")  # Ensure you upload this file to Hugging Face Spaces

import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.utils import load_img, img_to_array

# Load your trained model
model = tf.keras.models.load_model("denis_mnist_cnn_model_resnet50_v1.h5")  # Ensure you upload this file to Hugging Face Spaces 

# Define a function to preprocess the image
def preprocess_image(image):
    """
    Preprocesses the uploaded image for prediction.
    """
    image = image.resize((128, 128))  # Resize to match the model input size
    image = img_to_array(image)  # Convert PIL image to NumPy array
    image = preprocess_input(image)  # Normalize for ResNet50
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image

# Define the prediction function
def predict(image):
    """
    Accepts an image, preprocesses it, and returns the predicted label.
    """
    processed_image = preprocess_image(image)
    predictions = model.predict(processed_image)
    predicted_class = np.argmax(predictions, axis=-1)[0]  # Get the class index
    confidence = np.max(predictions)  # Get confidence score
    #return f"Predicted Class: {predicted_class}, Confidence: {confidence:.2f}"
    return {"prediction": int(predicted_class)}


# Create a Gradio interface
#interface = gr.Interface(fn=predict, inputs="image", outputs="json")


# Create a Gradio interface
interface = gr.Interface(
    fn=predict,  # The prediction function
    inputs=gr.Image(type="pil", label="Upload an Image"),  # Input: Image
    outputs=gr.Textbox(label="Prediction"),  # Output: Textbox
    title="MNIST ResNet50 Classifier",
    description="Upload an image to classify it using the trained ResNet50 model.",
    examples=[
        ["example_images/example1.png"],  # Add paths to example images in your Hugging Face repository
        ["example_images/example2.png"]
    ],
)


# Launch the app
if __name__ == "__main__":
    interface.launch(share=True)