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

# Carregar o modelo e processador
model = AutoModelForImageClassification.from_pretrained("./vit-finetuned")
processor = AutoProcessor.from_pretrained("google/vit-base-patch16-224-in21k")

# Função para prever a classe da imagem
def predict_image(image: Image.Image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted_class_idx = torch.argmax(logits, dim=-1).item()
    return predicted_class_idx

# Configuração do Streamlit
st.title("Modelo de Classificação de Imagens - Beans Dataset")
st.write("Carregue uma imagem de feijão para fazer uma predição!")

# Upload da imagem
uploaded_file = st.file_uploader("Escolha uma imagem", type=["jpg", "png", "jpeg"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    
    # Exibir a imagem
    st.image(image, caption="Imagem carregada", use_column_width=True)
    
    # Predição
    predicted_class_idx = predict_image(image)
    st.write(f"A classe predita é: {predicted_class_idx}")