class_img / app.py
KK-Sequeira's picture
Adiciona todos os arquivos na pasta
5bf3257
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}")