Update onnx_runtime_draw_bboxes.py

onnx_runtime_draw_bboxes.py

@@ -18,6 +18,17 @@ MODEL_INPUT_SIZE = (320, 320)  # Change based on your model input size
 # Define class labels (update based on your model)
 CLASS_NAMES = ["person", "car", "truck", "bicycle", "dog", "cat"]  # Update accordingly
 
+# Define unique colors for each class (BGR format for OpenCV)
+CLASS_COLORS = {
+    "person": (255, 0, 0),  # Blue
+    "car": (0, 255, 0),  # Green
+    "truck": (0, 0, 255),  # Red
+    "bicycle": (255, 255, 0),  # Cyan
+    "dog": (255, 0, 255),  # Magenta
+    "cat": (0, 255, 255),  # Yellow
+}
+DEFAULT_COLOR = (200, 200, 200)  # Gray for unknown classes
+
 def preprocess_image(image_path):
     """ Preprocess image for ONNX model input """
     image = cv2.imread(image_path)
@@ -44,15 +55,26 @@ def postprocess_output(output, orig_image):
             x1, y1, x2, y2 = boxes[i]  # Get box coordinates
             x1, y1, x2, y2 = int(x1 * width), int(y1 * height), int(x2 * width), int(y2 * height)  # Scale box
             
-            label = CLASS_NAMES[int(class_indices[i])]
+            # Get class label
+            class_id = int(class_indices[i])
+            label = CLASS_NAMES[class_id] if class_id < len(CLASS_NAMES) else "Unknown"
             confidence = scores[i]
 
+            # Get class-specific color
+            color = CLASS_COLORS.get(label, DEFAULT_COLOR)
+
             # Draw bounding box
-            cv2.rectangle(orig_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            cv2.rectangle(orig_image, (x1, y1), (x2, y2), color, 2)
 
             # Put label text
             label_text = f"{label}: {confidence:.2f}"
-            cv2.putText(orig_image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+            text_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
+            text_x, text_y = x1, y1 - 10
+
+            # Draw text background
+            cv2.rectangle(orig_image, (text_x, text_y - text_size[1] - 5), (text_x + text_size[0] + 5, text_y + 5), color, -1)
+            # Put label text
+            cv2.putText(orig_image, label_text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)
 
     return orig_image