Vannly's picture
Subida inicial del proyecto
6dd0b25
Raw
History Blame Contribute Delete
1.35 kB
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()