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

# Cargamos el modelo CLIP
clf = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")

def dog_cat_classifier(image):
    if image is None:
        return {"❌ Sin imagen": 1.0}
    
    image = image.convert("RGB")
    candidate_labels = ["dog", "cat", "neither dog nor cat"]
    preds = clf(image, candidate_labels=candidate_labels)
    
    result = {}
    for p in preds:
        label = p["label"]
        score = p["score"]
        if label == "dog":
            result["🐶 Dog"] = score
        elif label == "cat":
            result["🐱 Cat"] = score
        else:
            result["❌ Other"] = score
    return result

# --- DISEÑO DE LA INTERFAZ ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange", secondary_hue="slate")) as demo:
    
    # Encabezado con Estilo
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(
                """
                # 🐾 AI Pet Detector
                ### Clasificador inteligente basado en CLIP (Zero-Shot)
                Sube una foto y deja que la IA identifique si es un gato, un perro o algo totalmente distinto.
                """
            )
    
    gr.HTML("<hr>") # Una línea separadora elegante

    with gr.Row():
        # Columna Izquierda: Entrada
        with gr.Column(variant="panel"):
            gr.Markdown("### 📸 Paso 1: Sube tu imagen")
            input_img = gr.Image(type="pil", label=None)
            btn = gr.Button("🚀 Analizar Mascota", variant="primary")

        # Columna Derecha: Salida
        with gr.Column(variant="panel"):
            gr.Markdown("### 🐾 Paso 2: Resultado")
            output_label = gr.Label(num_top_classes=3, label="Probabilidades")
    
    # Ejemplos (Añaden mucha "presencia" visual)
    gr.Examples(
        examples=[["ejemplo_perro.jpg"], ["ejemplo_gato.jpg"]], # Asegúrate de tener estas imágenes o quita esta parte
        inputs=input_img
    )

    # Pie de página
    gr.Markdown(
        """
        ---
        *Desarrollado con ❤️ usando OpenAI CLIP y Gradio.*
        """
    )

    # Evento de clic
    btn.click(fn=dog_cat_classifier, inputs=input_img, outputs=output_label)

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