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

# Baixar os arquivos
repo_id = "davebraga/wrdbTI6"
model_path = hf_hub_download(repo_id, "trained_model.keras")
category_encoder_path = hf_hub_download(repo_id, "category_encoder.pkl")
color_encoder_path = hf_hub_download(repo_id, "color_encoder.pkl")

# Carregar modelo e encoders
model = load_model(model_path)
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)

# Previsão
def predict(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
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()