diff --git a/ultralytics/engine/results.py b/ultralytics/engine/results.py index 1c7ffb91..89b2d5c8 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, )