from utils import*
from PIL import Image

# dicionario das classes 
classes = {0: 'Avião', 1: 'Carro', 2: 'Pássaro', 3: 'Gato', 4: 'Cervo', 5: 'Cachorro', 6: 'Sapo', 7: 'Cavalo', 8: 'Navio', 9: 'Caminhão'}


import gradio as gr


model = Lenet5Modular()

model.load_state_dict(torch.load(arquivo))




# Função de classificação de imagem
def classifier_image(image):
    preprocess = create_transformer()

    # Converter a imagem para PIL se não estiver neste formato
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image.astype('uint8'), 'RGB')
    # Redimensionar a imagem para 32x32
    image = image.resize((32, 32))
    image_tensor = preprocess(image)
    image_tensor = image_tensor.unsqueeze(0)

    # Realizar a predição
    with torch.no_grad():
        output = model(image_tensor)
    _, predicted_class = torch.max(output, 1)

    return f"{str(classes[predicted_class.item()])}"

# Interface Gradio
gr.Interface(fn=classifier_image, inputs="image", outputs="text").launch()