import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# Load the model (ensure you have the correct model path)
model = tf.keras.models.load_model("denis_mnist_cnn_model.h5")

# Define a function to preprocess input and make predictions
def predict(image):
    # Convert image to a numpy array
    image = np.array(image)
    
    # Resize the image to the expected shape (28, 28, 3)
    image = tf.image.resize(image, (28, 28))  # Resize to 28x28 pixels
    image = np.expand_dims(image, axis=-1)  # Add the channel dimension if grayscale
    image = np.repeat(image, 3, axis=-1)  # Convert grayscale to RGB (if model was trained on RGB images)
    
    # Normalize the image
    image = image / 255.0
    
    # Add batch dimension
    image = np.expand_dims(image, axis=0)  # Add batch dimension to match the model's expected input shape (1, 28, 28, 3)
    
    # Perform prediction
    prediction = model.predict(image)
    
    # Return prediction as JSON
    return {"prediction": prediction.tolist()}

# Create a Gradio interface
interface = gr.Interface(
    fn=predict, 
    inputs="image",  # Image input for testing
    outputs="json"  # JSON output for prediction results
)

# Launch the interface
interface.launch(share=True)