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

# Cargamos un modelo ligero de clasificación
classifier = pipeline(
    "image-classification",
    model="umm-maybe/ai-image-detector",
)

def predict(image):
    # Convertir a formato PIL
    image = Image.fromarray(image)

    # Ejecutar predicción
    outputs = classifier(image)

    # Ordenamos por confianza
    outputs = sorted(outputs, key=lambda x: x["score"], reverse=True)
    
    label = outputs[0]["label"]
    score = outputs[0]["score"]

    # Formato de respuesta
    if "ai" in label.lower():
        result = f"🧠 Imagen generada por IA\nConfianza: {score:.2%}"
    else:
        result = f"📷 Imagen real\nConfianza: {score:.2%}"

    return result


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs="text",
    title="Detector de Imágenes IA",
    description="Sube una imagen y detecta si es real o generada por IA.",
)

demo.launch()