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

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

# Función para procesar la imagen y detectar objetos
def detect_objects(image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    # Ajustar tamaño de salida al de la imagen
    target_sizes = torch.tensor([image.size[::-1]])
    results = processor.post_process_object_detection(
        outputs, target_sizes=target_sizes, threshold=0.9
    )[0]

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

    return "\n".join(detected_objects)

# Crear interfaz con Gradio
def create_interface():
    interface = gr.Interface(
        fn=detect_objects,
        inputs=gr.Image(type="pil"),
        outputs=gr.Textbox(),
        live=True,
        title="Detección de Objetos con Transformers",
        description="Sube una imagen y descubre qué objetos puede detectar."
    )
    interface.launch()

if __name__ == "__main__":
    create_interface()