HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.2.77 new color_mode=instance plot arg (#15034)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Laughing 2024-08-14 00:45:53 +08:00 • committed by GitHub
parent b7c5db94b4
commit 265166ae4d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 32 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

24
ultralytics/engine/results.py
View file

@ -460,6 +460,7 @@ class Results(SimpleClass):
show=False,
save=False,
filename=None,
color_mode="class",
):
"""
Plots detection results on an input RGB image.
@ -481,6 +482,7 @@ class Results(SimpleClass):
show (bool): Whether to display the annotated image.
save (bool): Whether to save the annotated image.
filename (str | None): Filename to save image if save is True.
color_mode (bool): Specify the color mode, e.g., 'instance' or 'class'. Default to 'class'.
Returns:
(np.ndarray): Annotated image as a numpy array.
@ -491,6 +493,7 @@ class Results(SimpleClass):
... im = result.plot()
... im.show()
"""
assert color_mode in {"instance", "class"}, f"Expected color_mode='instance' or 'class', not {color_mode}."
if img is None and isinstance(self.orig_img, torch.Tensor):
img = (self.orig_img[0].detach().permute(1, 2, 0).contiguous() * 255).to(torch.uint8).cpu().numpy()
@ -519,17 +522,22 @@ class Results(SimpleClass):
.contiguous()
/ 255
)
idx = pred_boxes.cls if pred_boxes else range(len(pred_masks))
idx = 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
if pred_boxes is not None and show_boxes:
for d in reversed(pred_boxes):
for i, d in enumerate(reversed(pred_boxes)):
c, conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())
name = ("" if id is None else f"id:{id} ") + names[c]
label = (f"{name} {conf:.2f}" if conf else name) if labels else None
box = d.xyxyxyxy.reshape(-1, 4, 2).squeeze() if is_obb else d.xyxy.squeeze()
annotator.box_label(box, label, color=colors(c, True), rotated=is_obb)
annotator.box_label(
box,
label,
color=colors(i if color_mode == "instance" else c, True),
rotated=is_obb,
)
# Plot Classify results
if pred_probs is not None and show_probs:
@ -539,8 +547,14 @@ class Results(SimpleClass):
# Plot Pose results
if self.keypoints is not None:
for k in reversed(self.keypoints.data):
annotator.kpts(k, self.orig_shape, radius=kpt_radius, kpt_line=kpt_line)
for i, k in enumerate(reversed(self.keypoints.data)):
annotator.kpts(
k,
self.orig_shape,
radius=kpt_radius,
kpt_line=kpt_line,
kpt_color=colors(i, True) if color_mode == "instance" else None,
)
# Show results
if show: