import gradio as gr
from transformers import AutoImageProcessor, AutoModelForObjectDetection
import torch
from PIL import Image, ImageDraw

# Cargar modelo y procesador
checkpoint = "PablitoGil14/Practica2"
model = AutoModelForObjectDetection.from_pretrained(checkpoint)
processor = AutoImageProcessor.from_pretrained(checkpoint)

# Función de detección
def detectar_canguros(imagen: Image.Image):
    imagen = imagen.convert("RGB")
    inputs = processor(images=imagen, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = torch.tensor([imagen.size[::-1]]).to(model.device)
    results = processor.post_process(outputs, target_sizes=target_sizes)[0]

    draw = ImageDraw.Draw(imagen)
    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        if score >= 0.005:
            xmin, ymin, xmax, ymax = box.tolist()
            draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=3)
            draw.text((xmin + 4, ymin + 4), f"ID: {label.item()} ({round(score.item(), 2)})", fill="red")

    return imagen

# Interfaz Gradio
gr.Interface(
    fn=detectar_canguros,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Detector de Canguros 🦘",
    description="Sube una imagen y detecta canguros usando el modelo entrenado por PablitoGil14."
).launch()