import gradio as gr
from transformers import AutoModelForImageClassification, AutoImageProcessor, pipeline

# Carregar o modelo e o processador de imagens
model = AutoModelForImageClassification.from_pretrained("mestrevh/computer-vision-beans", use_safetensors=True)
image_processor = AutoImageProcessor.from_pretrained("mestrevh/computer-vision-beans")

# Criar o pipeline
classifier = pipeline("image-classification", model=model, feature_extractor=image_processor)

# Função de classificação
def predict_image(image):
    # A saída do classifier é uma lista de dicionários, pegar o label e a confiança
    result = classifier(image)
    label = result[0]['label']
    confidence = result[0]['score']
    return f"Class: {label}, Confidence: {confidence:.2f}"

# Interface Gradio
interface = gr.Interface(fn=predict_image, 
                         inputs=gr.Image(type="pil"), 
                         outputs="text", 
                         live=True)

interface.launch()