detr-finetuned

from transformers import AutoImageProcessor, AutoModelForObjectDetection
import torch
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt

# Abre la imagen desde una ruta en el equipo
url = "../pic/0fda72a2-f383-4f69-af8e-e16a0fbac621.jpg"
image = Image.open(url)

# Convierte la imagen a RGB si es gris
if image.mode != "RGB":
    image = image.convert("RGB")

processor = AutoImageProcessor.from_pretrained("0llheaven/detr-finetuned")
model = AutoModelForObjectDetection.from_pretrained("0llheaven/detr-finetuned")

inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)

# Filtra las predicciones con precisión mayor a 0.9
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)

print(results)
# Dibuja cuadros alrededor de los objetos detectados en la imagen
    draw = ImageDraw.Draw(image)
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()]}: {round(score.item(), 3)}", fill="red")

# Muestra la imagen
plt.figure(figsize=(10, 10))
plt.imshow(image)
plt.axis('off')
plt.show()