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

# Load your custom TensorFlow model. Update 'modelo_treinado.h5' with the path to your model.
tf_model_path = 'modelo_treinado.h5'
tf_model = tf.keras.models.load_model(tf_model_path)

# Define your class labels.
class_labels = ["Normal", "Cataract"]

def preprocess_image(image):
    # Resize the image to the input size required by the model (e.g., 224x224).
    image = image.resize((224, 224))
    # Convert the PIL image to a NumPy array and normalize pixel values.
    image = np.array(image) / 255.0
    # Add a batch dimension to the image.
    image = np.expand_dims(image, axis=0)
    return image

def predict(inp):
    # Preprocess the input image.
    inp = preprocess_image(inp)
    # Make predictions using your custom TensorFlow model.
    predictions = tf_model.predict(inp)
    # Get the class label with the highest confidence.
    predicted_class = class_labels[np.argmax(predictions)]
    # Get the confidence score of the predicted class.
    confidence = float(predictions[0][np.argmax(predictions)])
    
    # Create a dictionary with the predicted class and its confidence.
    result = {predicted_class: confidence}
    
    return result

# Create a Gradio interface.
gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(type="pil"),
    outputs=gr.outputs.Label(num_top_classes=1)
).launch()