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

# Carregar o modelo
print("Carregando o modelo...")
try:
    model = tf.keras.models.load_model("cnn_model.h5")
    print("Modelo carregado com sucesso.")
except Exception as e:
    print(f"Erro ao carregar o modelo: {e}")

# Função de predição
def predict(image):
    print("Imagem recebida.")  # Diagnóstico
    if image is None:
        print("Nenhuma imagem recebida.")
        return "Erro: Nenhuma imagem foi recebida."
    
    try:
        # Ajuste o tamanho da imagem conforme necessário
        image = image.resize((32, 32))
        image = np.array(image) / 255.0  # Normaliza a imagem
        image = np.expand_dims(image, axis=0)  # Expande as dimensões para o formato adequado
        print("Imagem processada para predição.")  # Diagnóstico

        # Realiza a predição
        prediction = model.predict(image)
        print(f"Predição realizada: {prediction}")  # Diagnóstico
        
        # Para classificação, pegamos a classe com maior probabilidade
        class_prediction = np.argmax(prediction, axis=-1)
        return class_prediction
    except Exception as e:
        print(f"Erro durante a predição: {e}")
        return f"Erro: {e}"

# Usando iface para definir a interface
iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Label())

# Lançando a interface
print("Lançando a interface...")
iface.launch()