import gradio as gr
from transformers import pipeline
from PIL import Image

# Carregar o modelo de classificação de feijão
classifier = pipeline("image-classification", model="mestrevh/computer-vision-beans")

# Função para classificar a imagem
def classify_bean(image):
    # Certifique-se de que a imagem é do tipo PIL
    if not isinstance(image, Image.Image):
        image = Image.open(image)
    
    # Realizar a classificação sem o padding, pois não é necessário para classificação de imagem
    result = classifier(image)
    
    # Retornar o resultado
    return result[0]['label'], result[0]['score']

# Criar a interface com Gradio
iface = gr.Interface(
    fn=classify_bean,  # Função que será chamada
    inputs=gr.Image(type="pil"),  # Usando PIL para garantir que a imagem esteja no formato correto
    outputs=["text", "json"],  # O que será exibido após a classificação
    live=True,  # Atualizar automaticamente
    title="Classificação de Feijões",  # Título do aplicativo
    description="Faça o upload de uma imagem de feijão e descubra a classificação!",  # Descrição
)

# Rodar o aplicativo
iface.launch(share=True)  # A opção share=True cria um link público