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

# Carregar o modelo TensorFlow
model = tf.keras.models.load_model('modelo_treinado.h5')

# Definir as classes
class_labels = ["Normal", "Cataract"]

# Função de previsão
def predict(inp):
    # Pré-processamento da imagem para adequá-la ao modelo TensorFlow
    img = np.array(inp)
    img = tf.image.resize(img, (224, 224))
    img = img / 255.0  # Normalização, se necessário
    img = tf.expand_dims(img, axis=0)

    # Fazer previsão com o modelo TensorFlow
    predictions = model.predict(img)
    
    # Obter a classe com a maior probabilidade
    predicted_class = class_labels[np.argmax(predictions)]
    
    return {predicted_class: float(predictions[0, np.argmax(predictions)])}

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