# app.py — versão FINAL, FUNCIONAL e PRECISA

from transformers import pipeline
from PIL import Image
import gradio as gr

# ✅ Modelo real e disponível: CLIP
classifier = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")

# Defina os rótulos em português (você controla!)
candidate_labels = [
    "uma foto de uma sala de estar",
    "uma foto de um quarto",
    "uma foto de uma cozinha",
    "uma foto de um banheiro",
    "uma foto de um escritório",
    "uma foto de uma sala de jantar",
    "uma foto de um corredor",
    "uma foto de um closet",
    "uma foto de uma garagem",
    "uma foto de uma varanda",
]

# Mapeamento para nomes curtos em PT-BR
label_map = {
    "uma foto de uma sala de estar": "Sala de estar",
    "uma foto de um quarto": "Quarto",
    "uma foto de uma cozinha": "Cozinha",
    "uma foto de um banheiro": "Banheiro",
    "uma foto de um escritório": "Escritório",
    "uma foto de uma sala de jantar": "Sala de jantar",
    "uma foto de um corredor": "Corredor",
    "uma foto de um closet": "Closet",
    "uma foto de uma garagem": "Garagem",
    "uma foto de uma varanda": "Varanda",
}

def classify_room(image):
    """Classifica o tipo de cômodo usando CLIP (zero-shot)."""
    if image is None:
        return "Nenhuma imagem fornecida", 0.0

    # CLIP espera PIL Image
    results = classifier(image, candidate_labels=candidate_labels, multi_label=False)

    # Pega o top-1
    top_result = results[0]
    label_en = top_result['label']
    confidence = top_result['score']
    label_pt = label_map.get(label_en, label_en)

    return label_pt, confidence

# Interface Gradio
demo = gr.Interface(
    fn=classify_room,
    inputs=gr.Image(type="pil", label="Envie uma foto do cômodo"),
    outputs=[
        gr.Textbox(label="Tipo de Cômodo Detectado"),
        gr.Number(label="Confiança", precision=4)
    ],
    title="🔍 Detector de Cômodos com CLIP (Zero-Shot)",
    description="Precisão muito melhorada — compara sua imagem com descrições em português.",
    # examples=[
    #    ["https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bedroom.jpg"],
    #   ["https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/living_room.jpg"]
    # ]
)

if __name__ == "__main__":
    demo.launch()