File size: 1,914 Bytes
834a10b 8a26bfc 29f2fd3 834a10b 8a26bfc 834a10b 29f2fd3 834a10b 29f2fd3 834a10b 29f2fd3 834a10b 8a26bfc 834a10b 29f2fd3 834a10b 29f2fd3 834a10b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import gradio as gr
import numpy as np
from tensorflow.keras.models import load_model
from huggingface_hub import hf_hub_download
import pickle
from PIL import Image
import zipfile
import os
# Pega o token do Hugging Face da variável de ambiente
hf_token = os.getenv("HUGGINGFACE_HUB_TOKEN")
# Repositório no Hugging Face
repo_id = "davebraga/wrdbTI6"
# Baixar o modelo compactado com autenticação
model_zip_path = hf_hub_download(repo_id, "saved_model.zip", token=hf_token)
# Diretório onde o modelo será extraído
extract_dir = "saved_model"
# Descompactar o modelo se ainda não existir
if not os.path.exists(extract_dir):
with zipfile.ZipFile(model_zip_path, 'r') as zip_ref:
zip_ref.extractall(extract_dir)
# Baixar os encoders com autenticação
category_encoder_path = hf_hub_download(repo_id, "category_encoder.pkl", token=hf_token)
color_encoder_path = hf_hub_download(repo_id, "color_encoder.pkl", token=hf_token)
# Carregar modelo e encoders
model = load_model(extract_dir)
with open(category_encoder_path, "rb") as f:
category_encoder = pickle.load(f)
with open(color_encoder_path, "rb") as f:
color_encoder = pickle.load(f)
# Função de previsão
def predict(image: Image.Image):
image = image.resize((160, 160))
image_array = np.array(image) / 255.0
image_array = np.expand_dims(image_array, axis=0)
category_pred, color_pred = model.predict(image_array)
category = category_encoder.inverse_transform([np.argmax(category_pred)])[0]
color = color_encoder.inverse_transform([np.argmax(color_pred)])[0]
return f"Categoria: {category}", f"Cor: {color}"
# Interface Gradio
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=["text", "text"],
title="Classificador de Categoria e Cor",
description="Faça upload de uma imagem de uma peça de roupa para prever a categoria e a cor."
)
iface.launch()
|