Update models/detection/detector.py

models/detection/detector.py

@@ -1,5 +1,6 @@
 import os
 import logging
+import random
 from PIL import Image, ImageDraw, ImageFont
 from huggingface_hub import hf_hub_download
 
@@ -67,36 +68,64 @@ class ObjectDetector:
         return detections
 
     def draw(self, image: Image.Image, detections, alpha=0.5):
+        """
+        Draws thicker, per-class-colored bounding boxes and labels.
+    
+        Args:
+            image (PIL.Image.Image): Original image.
+            detections (List[Dict]): Each dict has "bbox", "class_name", "confidence".
+            alpha (float): Blend strength for overlay.
+        Returns:
+            PIL.Image.Image: Blended image with overlays.
+        """
+        # copy & overlay
         overlay = image.copy()
         draw = ImageDraw.Draw(overlay)
     
-        # Load font
+        # try a TTF font, fallback to default
         try:
-            font = ImageFont.truetype("arial.ttf", 16)
+            font = ImageFont.truetype("arial.ttf", 18)
         except:
             font = ImageFont.load_default()
     
+        # deterministic color per class
+        class_colors = {}
+        def get_color(cls):
+            if cls not in class_colors:
+                # seed by class name → same color every run
+                rnd = random.Random(cls)
+                class_colors[cls] = (
+                    rnd.randint(100, 255),
+                    rnd.randint(100, 255),
+                    rnd.randint(100, 255),
+                )
+            return class_colors[cls]
+    
         for det in detections:
-            bbox = det["bbox"]
-            label = f'{det["class_name"]} {det["confidence"]:.2f}'
+            x1, y1, x2, y2 = det["bbox"]
+            cls_name = det["class_name"]
+            conf = det["confidence"]
+            label = f"{cls_name} {conf:.2f}"
+            color = get_color(cls_name)
     
-            # Draw thicker bounding box
-            for offset in range(3):
+            # thicker box: draw multiple offsets
+            for t in range(4):
                 draw.rectangle(
-                    [bbox[0] - offset, bbox[1] - offset, bbox[2] + offset, bbox[3] + offset],
-                    outline="red"
+                    (x1 - t, y1 - t, x2 + t, y2 + t),
+                    outline=color
                 )
     
-            # Calculate text size using textbbox
-            text_bbox = draw.textbbox((bbox[0], bbox[1]), label, font=font)
-            text_width = text_bbox[2] - text_bbox[0]
-            text_height = text_bbox[3] - text_bbox[1]
+            # calculate text size
+            text_box = draw.textbbox((x1, y1), label, font=font)
+            tb_w = text_box[2] - text_box[0]
+            tb_h = text_box[3] - text_box[1]
     
-            # Define background rectangle for text
-            text_bg = [bbox[0], bbox[1] - text_height, bbox[0] + text_width + 4, bbox[1]]
-            draw.rectangle(text_bg, fill="red")
+            # background rect for text
+            bg = (x1, y1 - tb_h, x1 + tb_w + 6, y1)
+            draw.rectangle(bg, fill=color)
     
-            # Draw text
-            draw.text((bbox[0] + 2, bbox[1] - text_height), label, fill="white", font=font)
+            # draw text (with small padding)
+            draw.text((x1 + 3, y1 - tb_h), label, font=font, fill="black")
     
-        return Image.blend(image, overlay, alpha)
+        # blend and return
+        return Image.blend(image, overlay, alpha)