Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import DetrImageProcessor, DetrForObjectDetection | |
import torch | |
from PIL import Image, ImageDraw | |
processor = DetrImageProcessor.from_pretrained("joortif/practica_2") | |
model = DetrForObjectDetection.from_pretrained("joortif/practica_2") | |
def detect_objects(img: Image.Image): | |
inputs = processor(images=img, return_tensors="pt") | |
outputs = model(**inputs) | |
target_sizes = torch.tensor([img.size[::-1]]) # | |
results = processor.post_process_object_detection( | |
outputs, target_sizes=target_sizes, threshold=0.9 | |
)[0] | |
draw = ImageDraw.Draw(img) | |
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): | |
box = [round(i, 2) for i in box.tolist()] | |
draw.rectangle(box, outline="red", width=3) | |
draw.text((box[0], box[1]), f"{model.config.id2label[label.item()]}: {score:.2f}", fill="red") | |
return img | |