from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import requests
import streamlit as st
import os
#from dotenv import load_dotenv

# Carregar variáveis de ambiente
#load_dotenv()

# Obter o token da API do Hugging Face
api_key = os.environ.get('API_KEY')  # Substitua pelo seu token

# Inicializar o pipeline com o modelo desejado
model_name = "openai/clip-vit-base-patch32"
processor = CLIPProcessor.from_pretrained(model_name, use_auth_token=api_key)
model = CLIPModel.from_pretrained(model_name, use_auth_token=api_key)

def classify_image(image, descriptions):
    inputs = processor(text=descriptions, images=image, return_tensors="pt", padding=True)
    outputs = model(**inputs)
    logits_per_image = outputs.logits_per_image
    probs = logits_per_image.softmax(dim=1)
    return probs

# Interface do Streamlit
st.title("Modelo de Classificação de Imagens")

# Carregar imagem
uploaded_file = st.file_uploader("Escolha uma imagem...(Cavalo, Cachorro, Criança, Gato)", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Imagem Carregada', use_column_width=True)
    
    # Definir descrições
    descriptions = ["uma foto de um gato", "uma foto de um cachorro", "uma foto de um cavalo", "uma foto de uma criança"]

    # Classificar imagem
    probs = classify_image(image, descriptions)

    # Mostrar resultados
    st.write("Resultados da classificação:")
    for description, prob in zip(descriptions, probs[0]):
        st.write(f"{description}: {prob:.4f}")