import streamlit as st
from transformers import AutoProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# Título do aplicativo
st.title("Classificador de Imagens com AutoTrain")

# Caminho do modelo no Hugging Face Model Hub (substitua pelo seu modelo AutoTrain)
model_name_or_path = "coan/designTrends"  # Substitua com seu modelo AutoTrain

# Carregar o processador e o modelo
@st.cache_resource
def load_model():
    processor = AutoProcessor.from_pretrained(model_name_or_path)
    model = AutoModelForImageClassification.from_pretrained(model_name_or_path)
    return processor, model

processor, model = load_model()

# Carregar a imagem pelo uploader
uploaded_file = st.file_uploader("Escolha uma imagem...", type=["jpeg", "jpg", "png"])

if uploaded_file is not None:
    # Abrir a imagem usando PIL
    image = Image.open(uploaded_file).convert("RGB")

    # Exibir a imagem no aplicativo
    st.image(image, caption='Imagem carregada.', use_column_width=True)

    # Pré-processar a imagem
    pixel_values = processor(images=image, return_tensors="pt").pixel_values

    # Fazer a previsão
    with torch.no_grad():
        outputs = model(pixel_values)
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()

        # Obter a label prevista
        predicted_label = model.config.id2label[predicted_class_idx]

    # Exibir o resultado da previsão
    st.write(f"Categoria prevista: **{predicted_label}**")