adjust visualisation style, allow draw_bbox() to show person id

rtmo_gpu.py

@@ -211,22 +211,24 @@ def draw_mmpose(img,
 
     return img
 
-def draw_bbox(img, bboxes, bboxes_scores=None, color=None):
+def draw_bbox(img, bboxes, bboxes_scores=None, color=None, person_id_list=None, line_width=2):
+    green = (0, 255, 0)
     for i, bbox in enumerate(bboxes):
         # Determine the color based on the score if no color is given
         if color is None and bboxes_scores is not None:
             # Scale the score to a color range (green to red)
             score = bboxes_scores[i]
-            green = int((1 - score) * 255)
-            red = int(score * 255)
-            box_color = (0, green, red)
+            start_color = np.array([128,128,128],dtype=np.uint8)
+            end_color = np.array([128,255,128],dtype=np.uint8)
+            box_color = (1 - score) * start_color + score * end_color
         else:
-            box_color = color if color is not None else (0, 255, 0)
+            box_color = color if color is not None else end_color
         
         # Draw the bounding box
         img = cv2.rectangle(img, (int(bbox[0]), int(bbox[1])),
-                            (int(bbox[2]), int(bbox[3])), box_color, 1)
+                            (int(bbox[2]), int(bbox[3])), box_color, line_width)
         
+        green_color = (0,255,0)
         # Display the score at the top-right corner of the bounding box
         if bboxes_scores is not None:
             score_text = f'{bboxes_scores[i]:.2f}'
@@ -235,6 +237,15 @@ def draw_bbox(img, bboxes, bboxes_scores=None, color=None):
             text_y = int(bbox[1]) + text_size[1]
             img = cv2.putText(img, score_text, (text_x, text_y),
                               cv2.FONT_HERSHEY_SIMPLEX, 0.5, box_color, 1, cv2.LINE_AA)
+
+        # Display Person ID on the top-right corner edge of the bounding box
+        if person_id_list is not None:
+            person_id_text = str(person_id_list[i])
+            text_size, _ = cv2.getTextSize(person_id_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
+            text_x = int(bbox[2]) - text_size[0]
+            text_y = int(bbox[1]) - text_size[1]
+            img = cv2.putText(img, person_id_text, (text_x, text_y),
+                              cv2.FONT_HERSHEY_SIMPLEX, 0.5, box_color, 2, cv2.LINE_AA)
     return img
 
 # with simplification to use onnxruntime only
@@ -242,7 +253,7 @@ def draw_skeleton(img,
                   keypoints,
                   scores,
                   kpt_thr=0.5,
-                  radius=2,
+                  radius=1,
                   line_width=2):
     num_keypoints = keypoints.shape[1]