adjust visualisation style, allow draw_bbox() to show person id
Browse files- rtmo_gpu.py +18 -7
rtmo_gpu.py
CHANGED
@@ -211,22 +211,24 @@ def draw_mmpose(img,
|
|
211 |
|
212 |
return img
|
213 |
|
214 |
-
def draw_bbox(img, bboxes, bboxes_scores=None, color=None):
|
|
|
215 |
for i, bbox in enumerate(bboxes):
|
216 |
# Determine the color based on the score if no color is given
|
217 |
if color is None and bboxes_scores is not None:
|
218 |
# Scale the score to a color range (green to red)
|
219 |
score = bboxes_scores[i]
|
220 |
-
|
221 |
-
|
222 |
-
box_color = (
|
223 |
else:
|
224 |
-
box_color = color if color is not None else
|
225 |
|
226 |
# Draw the bounding box
|
227 |
img = cv2.rectangle(img, (int(bbox[0]), int(bbox[1])),
|
228 |
-
(int(bbox[2]), int(bbox[3])), box_color,
|
229 |
|
|
|
230 |
# Display the score at the top-right corner of the bounding box
|
231 |
if bboxes_scores is not None:
|
232 |
score_text = f'{bboxes_scores[i]:.2f}'
|
@@ -235,6 +237,15 @@ def draw_bbox(img, bboxes, bboxes_scores=None, color=None):
|
|
235 |
text_y = int(bbox[1]) + text_size[1]
|
236 |
img = cv2.putText(img, score_text, (text_x, text_y),
|
237 |
cv2.FONT_HERSHEY_SIMPLEX, 0.5, box_color, 1, cv2.LINE_AA)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
238 |
return img
|
239 |
|
240 |
# with simplification to use onnxruntime only
|
@@ -242,7 +253,7 @@ def draw_skeleton(img,
|
|
242 |
keypoints,
|
243 |
scores,
|
244 |
kpt_thr=0.5,
|
245 |
-
radius=
|
246 |
line_width=2):
|
247 |
num_keypoints = keypoints.shape[1]
|
248 |
|
|
|
211 |
|
212 |
return img
|
213 |
|
214 |
+
def draw_bbox(img, bboxes, bboxes_scores=None, color=None, person_id_list=None, line_width=2):
|
215 |
+
green = (0, 255, 0)
|
216 |
for i, bbox in enumerate(bboxes):
|
217 |
# Determine the color based on the score if no color is given
|
218 |
if color is None and bboxes_scores is not None:
|
219 |
# Scale the score to a color range (green to red)
|
220 |
score = bboxes_scores[i]
|
221 |
+
start_color = np.array([128,128,128],dtype=np.uint8)
|
222 |
+
end_color = np.array([128,255,128],dtype=np.uint8)
|
223 |
+
box_color = (1 - score) * start_color + score * end_color
|
224 |
else:
|
225 |
+
box_color = color if color is not None else end_color
|
226 |
|
227 |
# Draw the bounding box
|
228 |
img = cv2.rectangle(img, (int(bbox[0]), int(bbox[1])),
|
229 |
+
(int(bbox[2]), int(bbox[3])), box_color, line_width)
|
230 |
|
231 |
+
green_color = (0,255,0)
|
232 |
# Display the score at the top-right corner of the bounding box
|
233 |
if bboxes_scores is not None:
|
234 |
score_text = f'{bboxes_scores[i]:.2f}'
|
|
|
237 |
text_y = int(bbox[1]) + text_size[1]
|
238 |
img = cv2.putText(img, score_text, (text_x, text_y),
|
239 |
cv2.FONT_HERSHEY_SIMPLEX, 0.5, box_color, 1, cv2.LINE_AA)
|
240 |
+
|
241 |
+
# Display Person ID on the top-right corner edge of the bounding box
|
242 |
+
if person_id_list is not None:
|
243 |
+
person_id_text = str(person_id_list[i])
|
244 |
+
text_size, _ = cv2.getTextSize(person_id_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
|
245 |
+
text_x = int(bbox[2]) - text_size[0]
|
246 |
+
text_y = int(bbox[1]) - text_size[1]
|
247 |
+
img = cv2.putText(img, person_id_text, (text_x, text_y),
|
248 |
+
cv2.FONT_HERSHEY_SIMPLEX, 0.5, box_color, 2, cv2.LINE_AA)
|
249 |
return img
|
250 |
|
251 |
# with simplification to use onnxruntime only
|
|
|
253 |
keypoints,
|
254 |
scores,
|
255 |
kpt_thr=0.5,
|
256 |
+
radius=1,
|
257 |
line_width=2):
|
258 |
num_keypoints = keypoints.shape[1]
|
259 |
|