Views
No views yet
1from PIL import Image
2import torch
3from transformers import YolosImageProcessor, YolosForObjectDetection
4
5device = 'cpu'
6if torch.cuda.is_available():
7 device = torch.device('cuda')
8elif torch.backends.mps.is_available():
9 device = torch.device('mps')
10
11ckpt = 'yainage90/fashion-object-detection-yolos-tiny'
12image_processor = YolosImageProcessor.from_pretrained(ckpt)
13model = YolosForObjectDetection.from_pretrained(ckpt).to(device)
14
15image = Image.open('<path/to/image>').convert('RGB')
16
17with torch.no_grad():
18 inputs = image_processor(images=[image], return_tensors="pt")
19 outputs = model(**inputs.to(device))
20 target_sizes = torch.tensor([[image.size[1], image.size[0]]])
21 results = image_processor.post_process_object_detection(outputs, threshold=0.85, target_sizes=target_sizes)[0]
22
23 items = []
24 for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
25 score = score.item()
26 label = label.item()
27 box = [i.item() for i in box]
28 print(f"{model.config.id2label[label]}: {round(score, 3)} at {box}")
29 items.append((score, label, box))