update README.md

README.md

@@ -70,6 +70,7 @@ pip install -r requirements.txt
 ```python
 from datasets import load_dataset
 from transformers import AutoProcessor, AutoModel
+from PIL import ImageDraw, ImageFont
 
 # Load the validation split (or 'train')
 ds = load_dataset("Laudando-Associates-LLC/pucks", split="test")
@@ -87,11 +88,21 @@ inputs = processor(image)
 # Run inference
 outputs = model(**inputs, conf_threshold=0.4)
 
-# Get outputs
+# Draw boxes
+draw = ImageDraw.Draw(image)
+font = ImageFont.truetype("DejaVuSans-Bold.ttf", size=24)
 for result in outputs:
-  boxes  = result["boxes"]   # [N, 4] bounding boxes (x1, y1, x2, y2)
-  labels = result["labels"]  # [N] class indices
-  scores = result["scores"]  # [N] confidence scores
+    boxes = result["boxes"]
+    labels = result["labels"]
+    scores = result["scores"]
+
+    for box, label, score in zip(boxes, labels, scores):
+        x1, y1, x2, y2 = box.tolist()
+        draw.rectangle([x1, y1, x2, y2], outline="blue", width=5)
+        draw.text((x1, max(0, y1 - 25)), f"{score:.2f}", fill="blue", font=font)
+
+# Save result
+image.save("output.jpg")
 ```
 
 ## How to Use