|
|
from ultralytics.utils.plotting import Annotator, colors |
|
|
import numpy as np |
|
|
import torch |
|
|
from copy import deepcopy |
|
|
|
|
|
def custom_plot( |
|
|
self, |
|
|
conf: bool = True, |
|
|
line_width: float | None = None, |
|
|
font_size: float | None = None, |
|
|
font: str = "Arial.ttf", |
|
|
pil: bool = False, |
|
|
img: np.ndarray | None = None, |
|
|
im_gpu: torch.Tensor | None = None, |
|
|
kpt_radius: int = 5, |
|
|
kpt_line: bool = True, |
|
|
labels: bool = True, |
|
|
boxes: bool = True, |
|
|
masks: bool = True, |
|
|
probs: bool = True, |
|
|
show: bool = False, |
|
|
save: bool = False, |
|
|
filename: str | None = None, |
|
|
color_mode: str = "class", |
|
|
txt_color: tuple[int, int, int] = (255, 255, 255), |
|
|
barcode_texts: list[str|None] = None, |
|
|
) -> np.ndarray: |
|
|
"""Plot detection results on an input BGR image. |
|
|
|
|
|
Args: |
|
|
conf (bool): Whether to plot detection confidence scores. |
|
|
line_width (float | None): Line width of bounding boxes. If None, scaled to image size. |
|
|
font_size (float | None): Font size for text. If None, scaled to image size. |
|
|
font (str): Font to use for text. |
|
|
pil (bool): Whether to return the image as a PIL Image. |
|
|
img (np.ndarray | None): Image to plot on. If None, uses original image. |
|
|
im_gpu (torch.Tensor | None): Normalized image on GPU for faster mask plotting. |
|
|
kpt_radius (int): Radius of drawn keypoints. |
|
|
kpt_line (bool): Whether to draw lines connecting keypoints. |
|
|
labels (bool): Whether to plot labels of bounding boxes. |
|
|
boxes (bool): Whether to plot bounding boxes. |
|
|
masks (bool): Whether to plot masks. |
|
|
probs (bool): Whether to plot classification probabilities. |
|
|
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 (str): Specify the color mode, e.g., 'instance' or 'class'. |
|
|
txt_color (tuple[int, int, int]): Text color in BGR format for classification output. |
|
|
|
|
|
Returns: |
|
|
(np.ndarray | PIL.Image.Image): Annotated image as a NumPy array (BGR) or PIL image (RGB) if `pil=True`. |
|
|
|
|
|
Examples: |
|
|
>>> results = model("image.jpg") |
|
|
>>> for result in results: |
|
|
>>> 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).byte().cpu().numpy() |
|
|
|
|
|
names = self.names |
|
|
is_obb = self.obb is not None |
|
|
pred_boxes, show_boxes = self.obb if is_obb else self.boxes, boxes |
|
|
pred_masks, show_masks = self.masks, masks |
|
|
pred_probs, show_probs = self.probs, probs |
|
|
annotator = Annotator( |
|
|
deepcopy(self.orig_img if img is None else img), |
|
|
line_width, |
|
|
font_size, |
|
|
font, |
|
|
pil or (pred_probs is not None and show_probs), |
|
|
example=names, |
|
|
) |
|
|
|
|
|
if pred_boxes is not None and show_boxes: |
|
|
for i, d in enumerate(reversed(pred_boxes)): |
|
|
c, d_conf, id = int(d.cls), float(d.conf) if conf else None, int(d.id.item()) if d.is_track else None |
|
|
name = ("" if id is None else f"id:{id} ") + names[c] |
|
|
if barcode_texts is None: |
|
|
label = (f"{name} {d_conf:.2f}" if conf else name) if labels else None |
|
|
else: |
|
|
label = barcode_texts[len(pred_boxes) - i - 1] |
|
|
|
|
|
box = d.xyxyxyxy.squeeze() if is_obb else d.xyxy.squeeze() |
|
|
annotator.box_label( |
|
|
box, |
|
|
label, |
|
|
color=colors(0 if not label else 6), |
|
|
) |
|
|
return annotator.result(pil) |
|
|
|
