Views
No views yet
hustvl/yolos-tiny architecture, optimized specifically for dense object detection on retail store shelves using the SKU-110k dataset layout.1from transformers import YolosForObjectDetection, YolosImageProcessor
2import torch
3
4device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
5
6model = YolosForObjectDetection.from_pretrained("Bartek7630/yolos-tiny-sku110k-refined").to(device)
7image_processor = YolosImageProcessor.from_pretrained("Bartek7630/yolos-tiny-sku110k-refined")
8
9### Full Inference & Visualization Script
10
11You can copy and run the complete pipeline below to test the model on any retail shelf image.
12
13```python
14import torch
15import requests
16from PIL import Image, ImageDraw
17from transformers import YolosForObjectDetection, YolosImageProcessor
18
19# 1. Configuration & Environment Setup
20device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
21model_id = "Bartek7630/yolos-tiny-sku110k-refined"
22
23# 2. Load Model and Processor directly from Hugging Face Hub
24model = YolosForObjectDetection.from_pretrained(model_id).to(device)
25image_processor = YolosImageProcessor.from_pretrained(model_id)
26model.eval()
27
28# 3. Load Input Image (Replace URL with local path if necessary)
29url = "[https://raw.githubusercontent.com/huggingface/transformers/main/tests/fixtures/tests_samples/COCO/000000039769.png](https://raw.githubusercontent.com/huggingface/transformers/main/tests/fixtures/tests_samples/COCO/000000039769.png)"
30image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
31
32# 4. Preprocessing
33inputs = image_processor(images=image, return_tensors="pt").to(device)
34
35# 5. Model Inference
36with torch.no_grad():
37 outputs = model(**inputs)
38
39# 6. Post-Processing (Bounding Box Denormalization)
40target_sizes = torch.tensor([image.size[::-1]]).to(device)
41results = image_processor.post_process_object_detection(outputs, threshold=0.3, target_sizes=target_sizes)[0]
42
43# 7. Visualization
44draw = ImageDraw.Draw(image)
45for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
46 box = [int(i) for i in box.tolist()]
47
48 # Draw red rectangles over detected SKU items
49 draw.rectangle(box, outline="red", width=3)
50
51# Save or display the finalized audit result
52image.save("shelf_audit_result.png")
53image.show()