Views
No views yet
1from transformers import DetrImageProcessor, DetrForObjectDetection
2import torch
3from PIL import Image, ImageOps
4import requests
5
6url = "https://github.com/Isalia20/DETR-finetune/blob/main/IMG_3507.jpg?raw=true"
7image = Image.open(requests.get(url, stream=True).raw)
8image = ImageOps.exif_transpose(image)
9
10# you can specify the revision tag if you don't want the timm dependency
11processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
12model = DetrForObjectDetection.from_pretrained("isalia99/detr-resnet-50-sku110k")
13model = model.eval()
14inputs = processor(images=image, return_tensors="pt")
15outputs = model(**inputs)
16
17# convert outputs (bounding boxes and class logits) to COCO API
18# let's only keep detections with score > 0.8
19target_sizes = torch.tensor([image.size[::-1]])
20results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.8)[0]
21
22for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
23 box = [round(i, 2) for i in box.tolist()]
24 print(
25 f"Detected {model.config.id2label[label.item()]} with confidence "
26 f"{round(score.item(), 3)} at location {box}"
27 )Detected LABEL_1 with confidence 0.983 at location [665.49, 480.05, 708.15, 650.11]
Detected LABEL_1 with confidence 0.938 at location [204.99, 1405.9, 239.9, 1546.5]
...
Detected LABEL_1 with confidence 0.998 at location [772.85, 169.49, 829.67, 372.18]
Detected LABEL_1 with confidence 0.999 at location [828.28, 1475.16, 874.37, 1593.43]