|
|
|
@ -10,11 +10,10 @@ from functools import lru_cache |
|
|
|
|
|
|
|
|
|
import numpy as np |
|
|
|
|
import torch |
|
|
|
|
import torchvision.transforms.functional as F |
|
|
|
|
|
|
|
|
|
from ultralytics.yolo.data.augment import LetterBox |
|
|
|
|
from ultralytics.yolo.utils import LOGGER, SimpleClass, deprecation_warn, ops |
|
|
|
|
from ultralytics.yolo.utils.plotting import Annotator, colors |
|
|
|
|
from ultralytics.yolo.utils.torch_utils import TORCHVISION_0_10 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class BaseTensor(SimpleClass): |
|
|
|
@ -160,6 +159,7 @@ class Results(SimpleClass): |
|
|
|
|
pil=False, |
|
|
|
|
example='abc', |
|
|
|
|
img=None, |
|
|
|
|
img_gpu=None, |
|
|
|
|
kpt_line=True, |
|
|
|
|
labels=True, |
|
|
|
|
boxes=True, |
|
|
|
@ -178,6 +178,7 @@ class Results(SimpleClass): |
|
|
|
|
pil (bool): Whether to return the image as a PIL Image. |
|
|
|
|
example (str): An example string to display. Useful for indicating the expected format of the output. |
|
|
|
|
img (numpy.ndarray): Plot to another image. if not, plot to original image. |
|
|
|
|
img_gpu (torch.Tensor): Normalized image in gpu with shape (1, 3, 640, 640), for faster mask plotting. |
|
|
|
|
kpt_line (bool): Whether to draw lines connecting keypoints. |
|
|
|
|
labels (bool): Whether to plot the label of bounding boxes. |
|
|
|
|
boxes (bool): Whether to plot the bounding boxes. |
|
|
|
@ -185,7 +186,7 @@ class Results(SimpleClass): |
|
|
|
|
probs (bool): Whether to plot classification probability |
|
|
|
|
|
|
|
|
|
Returns: |
|
|
|
|
(None) or (PIL.Image): If `pil` is True, a PIL Image is returned. Otherwise, nothing is returned. |
|
|
|
|
(numpy.ndarray): A numpy array of the annotated image. |
|
|
|
|
""" |
|
|
|
|
# Deprecation warn TODO: remove in 8.2 |
|
|
|
|
if 'show_conf' in kwargs: |
|
|
|
@ -200,6 +201,13 @@ class Results(SimpleClass): |
|
|
|
|
pred_probs, show_probs = self.probs, probs |
|
|
|
|
names = self.names |
|
|
|
|
keypoints = self.keypoints |
|
|
|
|
if pred_masks and show_masks: |
|
|
|
|
if img_gpu is None: |
|
|
|
|
img = LetterBox(pred_masks.shape[1:])(image=annotator.im) |
|
|
|
|
img_gpu = torch.as_tensor(img, dtype=torch.float16, device=pred_masks.masks.device).permute( |
|
|
|
|
2, 0, 1).flip(0).contiguous() / 255 |
|
|
|
|
annotator.masks(pred_masks.data, colors=[colors(x, True) for x in pred_boxes.cls], im_gpu=img_gpu) |
|
|
|
|
|
|
|
|
|
if pred_boxes and show_boxes: |
|
|
|
|
for d in 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()) |
|
|
|
@ -207,15 +215,6 @@ class Results(SimpleClass): |
|
|
|
|
label = (f'{name} {conf:.2f}' if conf else name) if labels else None |
|
|
|
|
annotator.box_label(d.xyxy.squeeze(), label, color=colors(c, True)) |
|
|
|
|
|
|
|
|
|
if pred_masks and show_masks: |
|
|
|
|
im = torch.as_tensor(annotator.im, dtype=torch.float16, device=pred_masks.data.device).permute(2, 0, |
|
|
|
|
1).flip(0) |
|
|
|
|
if TORCHVISION_0_10: |
|
|
|
|
im = F.resize(im.contiguous(), pred_masks.data.shape[1:], antialias=True) / 255 |
|
|
|
|
else: |
|
|
|
|
im = F.resize(im.contiguous(), pred_masks.data.shape[1:]) / 255 |
|
|
|
|
annotator.masks(pred_masks.data, colors=[colors(x, True) for x in pred_boxes.cls], im_gpu=im) |
|
|
|
|
|
|
|
|
|
if pred_probs is not None and show_probs: |
|
|
|
|
n5 = min(len(names), 5) |
|
|
|
|
top5i = pred_probs.argsort(0, descending=True)[:n5].tolist() # top 5 indices |
|
|
|
@ -226,7 +225,7 @@ class Results(SimpleClass): |
|
|
|
|
for k in reversed(keypoints): |
|
|
|
|
annotator.kpts(k, self.orig_shape, kpt_line=kpt_line) |
|
|
|
|
|
|
|
|
|
return np.asarray(annotator.im) if annotator.pil else annotator.im |
|
|
|
|
return annotator.result() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Boxes(BaseTensor): |
|
|
|
|