app.py

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

def detect_objects(image):

    # Preprocesamiento

    inputs = processor(images=image, return_tensors="pt")

 

    # Detectar objetos

    with torch.no_grad():

        outputs = model(**inputs)

 

    # Filtrar resultados

    target_sizes = torch.tensor([image.size[::-1]])  # (alto, ancho)

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

 

    # Crear una lista de los resultados con nombre y puntuación

    labels = results["labels"]

    scores = results["scores"]

    boxes = results["boxes"]

 

    # Mostrar los objetos detectados

    detected_objects = []

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

   

    return "\n".join(detected_objects)

# Crear la interfaz 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 se pueden detectar."

    )

   

    interface.launch()

 

if __name__ == "__main__":

    create_interface()