modified bbox display

inference.py

@@ -9,13 +9,61 @@ from models.segmenters.model_loader import load_segmenter
 from utils.video import extract_frames, write_video
 
 
-def draw_boxes(frame: np.ndarray, boxes: np.ndarray) -> np.ndarray:
+def _color_for_label(label: str) -> tuple[int, int, int]:
+    # Deterministic BGR color from label text.
+    value = abs(hash(label)) % 0xFFFFFF
+    blue = value & 0xFF
+    green = (value >> 8) & 0xFF
+    red = (value >> 16) & 0xFF
+    return (blue, green, red)
+
+
+def draw_boxes(
+    frame: np.ndarray,
+    boxes: np.ndarray,
+    labels: Optional[Sequence[int]] = None,
+    queries: Optional[Sequence[str]] = None,
+    label_names: Optional[Sequence[str]] = None,
+) -> np.ndarray:
     output = frame.copy()
     if boxes is None:
         return output
-    for box in boxes:
+    for idx, box in enumerate(boxes):
         x1, y1, x2, y2 = [int(coord) for coord in box]
-        cv2.rectangle(output, (x1, y1), (x2, y2), (0, 255, 0), thickness=2)
+        if label_names is not None and idx < len(label_names):
+            label = label_names[idx]
+        elif labels is not None and idx < len(labels) and queries is not None:
+            label_idx = int(labels[idx])
+            if 0 <= label_idx < len(queries):
+                label = queries[label_idx]
+            else:
+                label = f"label_{label_idx}"
+        else:
+            label = f"label_{idx}"
+        color = _color_for_label(label)
+        cv2.rectangle(output, (x1, y1), (x2, y2), color, thickness=2)
+        if label:
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            font_scale = 0.5
+            thickness = 1
+            text_size, baseline = cv2.getTextSize(label, font, font_scale, thickness)
+            text_w, text_h = text_size
+            pad = 4
+            text_x = x1
+            text_y = max(y1 - 6, text_h + pad)
+            box_top_left = (text_x, text_y - text_h - pad)
+            box_bottom_right = (text_x + text_w + pad, text_y + baseline)
+            cv2.rectangle(output, box_top_left, box_bottom_right, color, thickness=-1)
+            cv2.putText(
+                output,
+                label,
+                (text_x + pad // 2, text_y - 2),
+                font,
+                font_scale,
+                (255, 255, 255),
+                thickness,
+                lineType=cv2.LINE_AA,
+            )
     return output
 
 
@@ -103,7 +151,13 @@ def infer_frame(
     except Exception:
         logging.exception("Inference failed for queries %s", text_queries)
         raise
-    return draw_boxes(frame, result.boxes), detections
+    return draw_boxes(
+        frame,
+        result.boxes,
+        labels=result.labels,
+        queries=text_queries,
+        label_names=result.label_names,
+    ), detections
 
 
 def infer_segmentation_frame(