from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
from PIL import Image
import gradio as gr

# Cargar procesador y modelo preentrenado
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")

def detect_objects(image):
    inputs = processor(images=image, return_tensors="pt")

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

    target_sizes = torch.tensor([image.size[::-1]])

    results = processor.post_process_object_detection(
        outputs,
        target_sizes=target_sizes,
        threshold=0.9
    )[0]

    labels = results["labels"]
    scores = results["scores"]
    boxes = results["boxes"]

    detected_objects = []

    for score, label, box in zip(scores, labels, boxes):
        class_name = model.config.id2label[label.item()]
        detected_objects.append(
            f"Objeto: {class_name} | Score: {score:.3f} | Box: {box.tolist()}"
        )

    return "\n".join(detected_objects)

interface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(),
    title="Detección de Objetos con DETR (Transformer)",
    description="Sube una imagen y el modelo DETR detectará los objetos presentes."
)

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