From 7c61569630c032c78529da01a0934defbb815a17 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Wed, 11 Sep 2024 22:00:25 +0500 Subject: [PATCH] `bbox` and `masks` coloring based on `track-ids` (#16073) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com> Co-authored-by: Ultralytics Assistant <135830346+UltralyticsAssistant@users.noreply.github.com> --- ultralytics/engine/results.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/ultralytics/engine/results.py b/ultralytics/engine/results.py index 1c7ffb9144..89b2d5c8e1 100644 --- a/ultralytics/engine/results.py +++ b/ultralytics/engine/results.py @@ -522,7 +522,13 @@ class Results(SimpleClass): .contiguous() / 255 ) - idx = pred_boxes.cls if pred_boxes and color_mode == "class" else reversed(range(len(pred_masks))) + idx = ( + pred_boxes.id + if pred_boxes.id is not None and color_mode == "instance" + else pred_boxes.cls + if pred_boxes and color_mode == "class" + else reversed(range(len(pred_masks))) + ) annotator.masks(pred_masks.data, colors=[colors(x, True) for x in idx], im_gpu=im_gpu) # Plot Detect results @@ -535,7 +541,16 @@ class Results(SimpleClass): annotator.box_label( box, label, - color=colors(i if color_mode == "instance" else c, True), + color=colors( + c + if color_mode == "class" + else id + if id is not None + else i + if color_mode == "instance" + else None, + True, + ), rotated=is_obb, )