Update models/detection/detector.py

models/detection/detector.py

@@ -1,6 +1,6 @@
 import os
 import logging
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageFont
 from huggingface_hub import hf_hub_download
 
 logger = logging.getLogger(__name__)
@@ -69,9 +69,28 @@ class ObjectDetector:
     def draw(self, image: Image.Image, detections, alpha=0.5):
         overlay = image.copy()
         draw = ImageDraw.Draw(overlay)
+    
+        # Use a basic font (optional: can use truetype for better styling)
+        try:
+            font = ImageFont.truetype("arial.ttf", 16)  # Try a common system font
+        except:
+            font = ImageFont.load_default()
+    
         for det in detections:
             bbox = det["bbox"]
             label = f'{det["class_name"]} {det["confidence"]:.2f}'
-            draw.rectangle(bbox, outline="red", width=2)
-            draw.text((bbox[0], bbox[1]), label, fill="red")
-        return Image.blend(image, overlay, alpha)  
+    
+            # Thicker rectangle (3–4 px)
+            for offset in range(3):  # creates a thicker border
+                draw.rectangle(
+                    [bbox[0]-offset, bbox[1]-offset, bbox[2]+offset, bbox[3]+offset],
+                    outline="red"
+                )
+    
+            # Draw label with background box
+            text_size = draw.textsize(label, font=font)
+            text_bg = [bbox[0], bbox[1] - text_size[1], bbox[0] + text_size[0] + 4, bbox[1]]
+            draw.rectangle(text_bg, fill="red")
+            draw.text((bbox[0] + 2, bbox[1] - text_size[1]), label, fill="white", font=font)
+    
+        return Image.blend(image, overlay, alpha)