import gradio as gr  # Ensure you import gradio correctly
import tensorflow as tf
import numpy as np

# Load the Keras model
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) for RGB images
    image = tf.image.resize(image, (28, 28))  # Resize to 28x28 pixels
    
    # Check if the image is grayscale (single channel), and convert to RGB if necessary
    if image.shape[-1] == 1:  # If it's grayscale (single channel)
        image = np.repeat(image, 3, axis=-1)  # Convert grayscale to RGB by repeating the channel
    
    # 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)
    
    # Get the predicted class (index of the highest probability)
    predicted_class = np.argmax(prediction)
    
    # Return prediction as JSON (with the predicted class label)
    return {"prediction": int(predicted_class)}

# 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)