# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import math
import warnings
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union
import cv2
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from PIL import __version__ as pil_version
from ultralytics.utils import IS_COLAB, IS_KAGGLE, LOGGER, TryExcept, ops, plt_settings, threaded
from ultralytics.utils.checks import check_font, check_version, is_ascii
from ultralytics.utils.files import increment_path
class Colors:
"""
Ultralytics color palette for visualization and plotting.
This class provides methods to work with the Ultralytics color palette, including converting hex color codes to
RGB values and accessing predefined color schemes for object detection and pose estimation.
Attributes:
palette (List[tuple]): List of RGB color tuples for general use.
n (int): The number of colors in the palette.
pose_palette (np.ndarray): A specific color palette array for pose estimation with dtype np.uint8.
Examples:
>>> from ultralytics.utils.plotting import Colors
>>> colors = Colors()
>>> colors(5, True) # Returns BGR format: (221, 111, 255)
>>> colors(5, False) # Returns RGB format: (255, 111, 221)
## Ultralytics Color Palette
| Index | Color | HEX | RGB |
|-------|-------------------------------------------------------------------|-----------|-------------------|
| 0 | | `#042aff` | (4, 42, 255) |
| 1 | | `#0bdbeb` | (11, 219, 235) |
| 2 | | `#f3f3f3` | (243, 243, 243) |
| 3 | | `#00dfb7` | (0, 223, 183) |
| 4 | | `#111f68` | (17, 31, 104) |
| 5 | | `#ff6fdd` | (255, 111, 221) |
| 6 | | `#ff444f` | (255, 68, 79) |
| 7 | | `#cced00` | (204, 237, 0) |
| 8 | | `#00f344` | (0, 243, 68) |
| 9 | | `#bd00ff` | (189, 0, 255) |
| 10 | | `#00b4ff` | (0, 180, 255) |
| 11 | | `#dd00ba` | (221, 0, 186) |
| 12 | | `#00ffff` | (0, 255, 255) |
| 13 | | `#26c000` | (38, 192, 0) |
| 14 | | `#01ffb3` | (1, 255, 179) |
| 15 | | `#7d24ff` | (125, 36, 255) |
| 16 | | `#7b0068` | (123, 0, 104) |
| 17 | | `#ff1b6c` | (255, 27, 108) |
| 18 | | `#fc6d2f` | (252, 109, 47) |
| 19 | | `#a2ff0b` | (162, 255, 11) |
## Pose Color Palette
| Index | Color | HEX | RGB |
|-------|-------------------------------------------------------------------|-----------|-------------------|
| 0 | | `#ff8000` | (255, 128, 0) |
| 1 | | `#ff9933` | (255, 153, 51) |
| 2 | | `#ffb266` | (255, 178, 102) |
| 3 | | `#e6e600` | (230, 230, 0) |
| 4 | | `#ff99ff` | (255, 153, 255) |
| 5 | | `#99ccff` | (153, 204, 255) |
| 6 | | `#ff66ff` | (255, 102, 255) |
| 7 | | `#ff33ff` | (255, 51, 255) |
| 8 | | `#66b2ff` | (102, 178, 255) |
| 9 | | `#3399ff` | (51, 153, 255) |
| 10 | | `#ff9999` | (255, 153, 153) |
| 11 | | `#ff6666` | (255, 102, 102) |
| 12 | | `#ff3333` | (255, 51, 51) |
| 13 | | `#99ff99` | (153, 255, 153) |
| 14 | | `#66ff66` | (102, 255, 102) |
| 15 | | `#33ff33` | (51, 255, 51) |
| 16 | | `#00ff00` | (0, 255, 0) |
| 17 | | `#0000ff` | (0, 0, 255) |
| 18 | | `#ff0000` | (255, 0, 0) |
| 19 | | `#ffffff` | (255, 255, 255) |
!!! note "Ultralytics Brand Colors"
For Ultralytics brand colors see [https://www.ultralytics.com/brand](https://www.ultralytics.com/brand).
Please use the official Ultralytics colors for all marketing materials.
"""
def __init__(self):
"""Initialize colors as hex = matplotlib.colors.TABLEAU_COLORS.values()."""
hexs = (
"042AFF",
"0BDBEB",
"F3F3F3",
"00DFB7",
"111F68",
"FF6FDD",
"FF444F",
"CCED00",
"00F344",
"BD00FF",
"00B4FF",
"DD00BA",
"00FFFF",
"26C000",
"01FFB3",
"7D24FF",
"7B0068",
"FF1B6C",
"FC6D2F",
"A2FF0B",
)
self.palette = [self.hex2rgb(f"#{c}") for c in hexs]
self.n = len(self.palette)
self.pose_palette = np.array(
[
[255, 128, 0],
[255, 153, 51],
[255, 178, 102],
[230, 230, 0],
[255, 153, 255],
[153, 204, 255],
[255, 102, 255],
[255, 51, 255],
[102, 178, 255],
[51, 153, 255],
[255, 153, 153],
[255, 102, 102],
[255, 51, 51],
[153, 255, 153],
[102, 255, 102],
[51, 255, 51],
[0, 255, 0],
[0, 0, 255],
[255, 0, 0],
[255, 255, 255],
],
dtype=np.uint8,
)
def __call__(self, i: int, bgr: bool = False) -> tuple:
"""
Convert hex color codes to RGB values.
Args:
i (int): Color index.
bgr (bool, optional): Whether to return BGR format instead of RGB.
Returns:
(tuple): RGB or BGR color tuple.
"""
c = self.palette[int(i) % self.n]
return (c[2], c[1], c[0]) if bgr else c
@staticmethod
def hex2rgb(h: str) -> tuple:
"""Convert hex color codes to RGB values (i.e. default PIL order)."""
return tuple(int(h[1 + i : 1 + i + 2], 16) for i in (0, 2, 4))
colors = Colors() # create instance for 'from utils.plots import colors'
class Annotator:
"""
Ultralytics Annotator for train/val mosaics and JPGs and predictions annotations.
Attributes:
im (Image.Image | np.ndarray): The image to annotate.
pil (bool): Whether to use PIL or cv2 for drawing annotations.
font (ImageFont.truetype | ImageFont.load_default): Font used for text annotations.
lw (float): Line width for drawing.
skeleton (List[List[int]]): Skeleton structure for keypoints.
limb_color (List[int]): Color palette for limbs.
kpt_color (List[int]): Color palette for keypoints.
dark_colors (set): Set of colors considered dark for text contrast.
light_colors (set): Set of colors considered light for text contrast.
Examples:
>>> from ultralytics.utils.plotting import Annotator
>>> im0 = cv2.imread("test.png")
>>> annotator = Annotator(im0, line_width=10)
>>> annotator.box_label([10, 10, 100, 100], "person", (255, 0, 0))
"""
def __init__(
self,
im,
line_width: Optional[int] = None,
font_size: Optional[int] = None,
font: str = "Arial.ttf",
pil: bool = False,
example: str = "abc",
):
"""Initialize the Annotator class with image and line width along with color palette for keypoints and limbs."""
non_ascii = not is_ascii(example) # non-latin labels, i.e. asian, arabic, cyrillic
input_is_pil = isinstance(im, Image.Image)
self.pil = pil or non_ascii or input_is_pil
self.lw = line_width or max(round(sum(im.size if input_is_pil else im.shape) / 2 * 0.003), 2)
if not input_is_pil:
if im.shape[2] == 1: # handle grayscale
im = cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)
elif im.shape[2] > 3: # multispectral
im = np.ascontiguousarray(im[..., :3])
if self.pil: # use PIL
self.im = im if input_is_pil else Image.fromarray(im)
if self.im.mode not in {"RGB", "RGBA"}: # multispectral
self.im = self.im.convert("RGB")
self.draw = ImageDraw.Draw(self.im, "RGBA")
try:
font = check_font("Arial.Unicode.ttf" if non_ascii else font)
size = font_size or max(round(sum(self.im.size) / 2 * 0.035), 12)
self.font = ImageFont.truetype(str(font), size)
except Exception:
self.font = ImageFont.load_default()
# Deprecation fix for w, h = getsize(string) -> _, _, w, h = getbox(string)
if check_version(pil_version, "9.2.0"):
self.font.getsize = lambda x: self.font.getbbox(x)[2:4] # text width, height
else: # use cv2
assert im.data.contiguous, "Image not contiguous. Apply np.ascontiguousarray(im) to Annotator input images."
self.im = im if im.flags.writeable else im.copy()
self.tf = max(self.lw - 1, 1) # font thickness
self.sf = self.lw / 3 # font scale
# Pose
self.skeleton = [
[16, 14],
[14, 12],
[17, 15],
[15, 13],
[12, 13],
[6, 12],
[7, 13],
[6, 7],
[6, 8],
[7, 9],
[8, 10],
[9, 11],
[2, 3],
[1, 2],
[1, 3],
[2, 4],
[3, 5],
[4, 6],
[5, 7],
]
self.limb_color = colors.pose_palette[[9, 9, 9, 9, 7, 7, 7, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]]
self.kpt_color = colors.pose_palette[[16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]]
self.dark_colors = {
(235, 219, 11),
(243, 243, 243),
(183, 223, 0),
(221, 111, 255),
(0, 237, 204),
(68, 243, 0),
(255, 255, 0),
(179, 255, 1),
(11, 255, 162),
}
self.light_colors = {
(255, 42, 4),
(79, 68, 255),
(255, 0, 189),
(255, 180, 0),
(186, 0, 221),
(0, 192, 38),
(255, 36, 125),
(104, 0, 123),
(108, 27, 255),
(47, 109, 252),
(104, 31, 17),
}
def get_txt_color(self, color: tuple = (128, 128, 128), txt_color: tuple = (255, 255, 255)) -> tuple:
"""
Assign text color based on background color.
Args:
color (tuple, optional): The background color of the rectangle for text (B, G, R).
txt_color (tuple, optional): The color of the text (R, G, B).
Returns:
(tuple): Text color for label.
Examples:
>>> from ultralytics.utils.plotting import Annotator
>>> im0 = cv2.imread("test.png")
>>> annotator = Annotator(im0, line_width=10)
>>> annotator.get_txt_color(color=(104, 31, 17)) # return (255, 255, 255)
"""
if color in self.dark_colors:
return 104, 31, 17
elif color in self.light_colors:
return 255, 255, 255
else:
return txt_color
def box_label(self, box, label: str = "", color: tuple = (128, 128, 128), txt_color: tuple = (255, 255, 255)):
"""
Draw a bounding box on an image with a given label.
Args:
box (tuple): The bounding box coordinates (x1, y1, x2, y2).
label (str, optional): The text label to be displayed.
color (tuple, optional): The background color of the rectangle (B, G, R).
txt_color (tuple, optional): The color of the text (R, G, B).
Examples:
>>> from ultralytics.utils.plotting import Annotator
>>> im0 = cv2.imread("test.png")
>>> annotator = Annotator(im0, line_width=10)
>>> annotator.box_label(box=[10, 20, 30, 40], label="person")
"""
txt_color = self.get_txt_color(color, txt_color)
if isinstance(box, torch.Tensor):
box = box.tolist()
multi_points = isinstance(box[0], list) # multiple points with shape (n, 2)
p1 = [int(b) for b in box[0]] if multi_points else (int(box[0]), int(box[1]))
if self.pil:
self.draw.polygon(
[tuple(b) for b in box], width=self.lw, outline=color
) if multi_points else self.draw.rectangle(box, width=self.lw, outline=color)
if label:
w, h = self.font.getsize(label) # text width, height
outside = p1[1] >= h # label fits outside box
if p1[0] > self.im.size[0] - w: # size is (w, h), check if label extend beyond right side of image
p1 = self.im.size[0] - w, p1[1]
self.draw.rectangle(
(p1[0], p1[1] - h if outside else p1[1], p1[0] + w + 1, p1[1] + 1 if outside else p1[1] + h + 1),
fill=color,
)
# self.draw.text([box[0], box[1]], label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0
self.draw.text((p1[0], p1[1] - h if outside else p1[1]), label, fill=txt_color, font=self.font)
else: # cv2
cv2.polylines(
self.im, [np.asarray(box, dtype=int)], True, color, self.lw
) if multi_points else cv2.rectangle(
self.im, p1, (int(box[2]), int(box[3])), color, thickness=self.lw, lineType=cv2.LINE_AA
)
if label:
w, h = cv2.getTextSize(label, 0, fontScale=self.sf, thickness=self.tf)[0] # text width, height
h += 3 # add pixels to pad text
outside = p1[1] >= h # label fits outside box
if p1[0] > self.im.shape[1] - w: # shape is (h, w), check if label extend beyond right side of image
p1 = self.im.shape[1] - w, p1[1]
p2 = p1[0] + w, p1[1] - h if outside else p1[1] + h
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
cv2.putText(
self.im,
label,
(p1[0], p1[1] - 2 if outside else p1[1] + h - 1),
0,
self.sf,
txt_color,
thickness=self.tf,
lineType=cv2.LINE_AA,
)
def masks(self, masks, colors, im_gpu, alpha: float = 0.5, retina_masks: bool = False):
"""
Plot masks on image.
Args:
masks (torch.Tensor): Predicted masks on cuda, shape: [n, h, w]
colors (List[List[int]]): Colors for predicted masks, [[r, g, b] * n]
im_gpu (torch.Tensor): Image is in cuda, shape: [3, h, w], range: [0, 1]
alpha (float, optional): Mask transparency: 0.0 fully transparent, 1.0 opaque.
retina_masks (bool, optional): Whether to use high resolution masks or not.
"""
if self.pil:
# Convert to numpy first
self.im = np.asarray(self.im).copy()
if len(masks) == 0:
self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
if im_gpu.device != masks.device:
im_gpu = im_gpu.to(masks.device)
colors = torch.tensor(colors, device=masks.device, dtype=torch.float32) / 255.0 # shape(n,3)
colors = colors[:, None, None] # shape(n,1,1,3)
masks = masks.unsqueeze(3) # shape(n,h,w,1)
masks_color = masks * (colors * alpha) # shape(n,h,w,3)
inv_alpha_masks = (1 - masks * alpha).cumprod(0) # shape(n,h,w,1)
mcs = masks_color.max(dim=0).values # shape(n,h,w,3)
im_gpu = im_gpu.flip(dims=[0]) # flip channel
im_gpu = im_gpu.permute(1, 2, 0).contiguous() # shape(h,w,3)
im_gpu = im_gpu * inv_alpha_masks[-1] + mcs
im_mask = im_gpu * 255
im_mask_np = im_mask.byte().cpu().numpy()
self.im[:] = im_mask_np if retina_masks else ops.scale_image(im_mask_np, self.im.shape)
if self.pil:
# Convert im back to PIL and update draw
self.fromarray(self.im)
def kpts(
self,
kpts,
shape: tuple = (640, 640),
radius: Optional[int] = None,
kpt_line: bool = True,
conf_thres: float = 0.25,
kpt_color: Optional[tuple] = None,
):
"""
Plot keypoints on the image.
Args:
kpts (torch.Tensor): Keypoints, shape [17, 3] (x, y, confidence).
shape (tuple, optional): Image shape (h, w).
radius (int, optional): Keypoint radius.
kpt_line (bool, optional): Draw lines between keypoints.
conf_thres (float, optional): Confidence threshold.
kpt_color (tuple, optional): Keypoint color (B, G, R).
Note:
- `kpt_line=True` currently only supports human pose plotting.
- Modifies self.im in-place.
- If self.pil is True, converts image to numpy array and back to PIL.
"""
radius = radius if radius is not None else self.lw
if self.pil:
# Convert to numpy first
self.im = np.asarray(self.im).copy()
nkpt, ndim = kpts.shape
is_pose = nkpt == 17 and ndim in {2, 3}
kpt_line &= is_pose # `kpt_line=True` for now only supports human pose plotting
for i, k in enumerate(kpts):
color_k = kpt_color or (self.kpt_color[i].tolist() if is_pose else colors(i))
x_coord, y_coord = k[0], k[1]
if x_coord % shape[1] != 0 and y_coord % shape[0] != 0:
if len(k) == 3:
conf = k[2]
if conf < conf_thres:
continue
cv2.circle(self.im, (int(x_coord), int(y_coord)), radius, color_k, -1, lineType=cv2.LINE_AA)
if kpt_line:
ndim = kpts.shape[-1]
for i, sk in enumerate(self.skeleton):
pos1 = (int(kpts[(sk[0] - 1), 0]), int(kpts[(sk[0] - 1), 1]))
pos2 = (int(kpts[(sk[1] - 1), 0]), int(kpts[(sk[1] - 1), 1]))
if ndim == 3:
conf1 = kpts[(sk[0] - 1), 2]
conf2 = kpts[(sk[1] - 1), 2]
if conf1 < conf_thres or conf2 < conf_thres:
continue
if pos1[0] % shape[1] == 0 or pos1[1] % shape[0] == 0 or pos1[0] < 0 or pos1[1] < 0:
continue
if pos2[0] % shape[1] == 0 or pos2[1] % shape[0] == 0 or pos2[0] < 0 or pos2[1] < 0:
continue
cv2.line(
self.im,
pos1,
pos2,
kpt_color or self.limb_color[i].tolist(),
thickness=int(np.ceil(self.lw / 2)),
lineType=cv2.LINE_AA,
)
if self.pil:
# Convert im back to PIL and update draw
self.fromarray(self.im)
def rectangle(self, xy, fill=None, outline=None, width: int = 1):
"""Add rectangle to image (PIL-only)."""
self.draw.rectangle(xy, fill, outline, width)
def text(self, xy, text: str, txt_color: tuple = (255, 255, 255), anchor: str = "top", box_color: tuple = ()):
"""
Add text to an image using PIL or cv2.
Args:
xy (List[int]): Top-left coordinates for text placement.
text (str): Text to be drawn.
txt_color (tuple, optional): Text color (R, G, B).
anchor (str, optional): Text anchor position ('top' or 'bottom').
box_color (tuple, optional): Box color (R, G, B, A) with optional alpha.
"""
if self.pil:
w, h = self.font.getsize(text)
if anchor == "bottom": # start y from font bottom
xy[1] += 1 - h
for line in text.split("\n"):
if box_color:
# Draw rectangle for each line
w, h = self.font.getsize(line)
self.draw.rectangle((xy[0], xy[1], xy[0] + w + 1, xy[1] + h + 1), fill=box_color)
self.draw.text(xy, line, fill=txt_color, font=self.font)
xy[1] += h
else:
if box_color:
w, h = cv2.getTextSize(text, 0, fontScale=self.sf, thickness=self.tf)[0]
h += 3 # add pixels to pad text
outside = xy[1] >= h # label fits outside box
p2 = xy[0] + w, xy[1] - h if outside else xy[1] + h
cv2.rectangle(self.im, xy, p2, box_color, -1, cv2.LINE_AA) # filled
cv2.putText(self.im, text, xy, 0, self.sf, txt_color, thickness=self.tf, lineType=cv2.LINE_AA)
def fromarray(self, im):
"""Update self.im from a numpy array."""
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
self.draw = ImageDraw.Draw(self.im)
def result(self):
"""Return annotated image as array."""
return np.asarray(self.im)
def show(self, title: Optional[str] = None):
"""Show the annotated image."""
im = Image.fromarray(np.asarray(self.im)[..., ::-1]) # Convert numpy array to PIL Image with RGB to BGR
if IS_COLAB or IS_KAGGLE: # can not use IS_JUPYTER as will run for all ipython environments
try:
display(im) # noqa - display() function only available in ipython environments
except ImportError as e:
LOGGER.warning(f"Unable to display image in Jupyter notebooks: {e}")
else:
im.show(title=title)
def save(self, filename: str = "image.jpg"):
"""Save the annotated image to 'filename'."""
cv2.imwrite(filename, np.asarray(self.im))
@staticmethod
def get_bbox_dimension(bbox: Optional[tuple] = None):
"""
Calculate the dimensions and area of a bounding box.
Args:
bbox (tuple): Bounding box coordinates in the format (x_min, y_min, x_max, y_max).
Returns:
width (float): Width of the bounding box.
height (float): Height of the bounding box.
area (float): Area enclosed by the bounding box.
Examples:
>>> from ultralytics.utils.plotting import Annotator
>>> im0 = cv2.imread("test.png")
>>> annotator = Annotator(im0, line_width=10)
>>> annotator.get_bbox_dimension(bbox=[10, 20, 30, 40])
"""
x_min, y_min, x_max, y_max = bbox
width = x_max - x_min
height = y_max - y_min
return width, height, width * height
@TryExcept() # known issue https://github.com/ultralytics/yolov5/issues/5395
@plt_settings()
def plot_labels(boxes, cls, names=(), save_dir=Path(""), on_plot=None):
"""
Plot training labels including class histograms and box statistics.
Args:
boxes (np.ndarray): Bounding box coordinates in format [x, y, width, height].
cls (np.ndarray): Class indices.
names (dict, optional): Dictionary mapping class indices to class names.
save_dir (Path, optional): Directory to save the plot.
on_plot (Callable, optional): Function to call after plot is saved.
"""
import matplotlib.pyplot as plt # scope for faster 'import ultralytics'
import pandas
from matplotlib.colors import LinearSegmentedColormap
# Filter matplotlib>=3.7.2 warning
warnings.filterwarnings("ignore", category=UserWarning, message="The figure layout has changed to tight")
warnings.filterwarnings("ignore", category=FutureWarning)
# Plot dataset labels
LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ")
nc = int(cls.max() + 1) # number of classes
boxes = boxes[:1000000] # limit to 1M boxes
x = pandas.DataFrame(boxes, columns=["x", "y", "width", "height"])
try: # Seaborn correlogram
import seaborn
seaborn.pairplot(x, corner=True, diag_kind="auto", kind="hist", diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
plt.savefig(save_dir / "labels_correlogram.jpg", dpi=200)
plt.close()
except ImportError:
pass # Skip if seaborn is not installed
# Matplotlib labels
subplot_3_4_color = LinearSegmentedColormap.from_list("white_blue", ["white", "blue"])
ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
y = ax[0].hist(cls, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
for i in range(nc):
y[2].patches[i].set_color([x / 255 for x in colors(i)])
ax[0].set_ylabel("instances")
if 0 < len(names) < 30:
ax[0].set_xticks(range(len(names)))
ax[0].set_xticklabels(list(names.values()), rotation=90, fontsize=10)
else:
ax[0].set_xlabel("classes")
boxes = np.column_stack([0.5 - boxes[:, 2:4] / 2, 0.5 + boxes[:, 2:4] / 2]) * 1000
img = Image.fromarray(np.ones((1000, 1000, 3), dtype=np.uint8) * 255)
for cls, box in zip(cls[:500], boxes[:500]):
ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls)) # plot
ax[1].imshow(img)
ax[1].axis("off")
ax[2].hist2d(x["x"], x["y"], bins=50, cmap=subplot_3_4_color)
ax[2].set_xlabel("x")
ax[2].set_ylabel("y")
ax[3].hist2d(x["width"], x["height"], bins=50, cmap=subplot_3_4_color)
ax[3].set_xlabel("width")
ax[3].set_ylabel("height")
for a in {0, 1, 2, 3}:
for s in {"top", "right", "left", "bottom"}:
ax[a].spines[s].set_visible(False)
fname = save_dir / "labels.jpg"
plt.savefig(fname, dpi=200)
plt.close()
if on_plot:
on_plot(fname)
def save_one_box(
xyxy,
im,
file: Path = Path("im.jpg"),
gain: float = 1.02,
pad: int = 10,
square: bool = False,
BGR: bool = False,
save: bool = True,
):
"""
Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop.
This function takes a bounding box and an image, and then saves a cropped portion of the image according
to the bounding box. Optionally, the crop can be squared, and the function allows for gain and padding
adjustments to the bounding box.
Args:
xyxy (torch.Tensor | list): A tensor or list representing the bounding box in xyxy format.
im (np.ndarray): The input image.
file (Path, optional): The path where the cropped image will be saved.
gain (float, optional): A multiplicative factor to increase the size of the bounding box.
pad (int, optional): The number of pixels to add to the width and height of the bounding box.
square (bool, optional): If True, the bounding box will be transformed into a square.
BGR (bool, optional): If True, the image will be returned in BGR format, otherwise in RGB.
save (bool, optional): If True, the cropped image will be saved to disk.
Returns:
(np.ndarray): The cropped image.
Examples:
>>> from ultralytics.utils.plotting import save_one_box
>>> xyxy = [50, 50, 150, 150]
>>> im = cv2.imread("image.jpg")
>>> cropped_im = save_one_box(xyxy, im, file="cropped.jpg", square=True)
"""
if not isinstance(xyxy, torch.Tensor): # may be list
xyxy = torch.stack(xyxy)
b = ops.xyxy2xywh(xyxy.view(-1, 4)) # boxes
if square:
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
xyxy = ops.xywh2xyxy(b).long()
xyxy = ops.clip_boxes(xyxy, im.shape)
grayscale = im.shape[2] == 1 # grayscale image
crop = im[int(xyxy[0, 1]) : int(xyxy[0, 3]), int(xyxy[0, 0]) : int(xyxy[0, 2]), :: (1 if BGR or grayscale else -1)]
if save:
file.parent.mkdir(parents=True, exist_ok=True) # make directory
f = str(increment_path(file).with_suffix(".jpg"))
# cv2.imwrite(f, crop) # save BGR, https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
crop = crop.squeeze(-1) if grayscale else crop[..., ::-1] if BGR else crop
Image.fromarray(crop).save(f, quality=95, subsampling=0) # save RGB
return crop
@threaded
def plot_images(
labels: Dict[str, Any],
images: Union[torch.Tensor, np.ndarray] = np.zeros((0, 3, 640, 640), dtype=np.float32),
paths: Optional[List[str]] = None,
fname: str = "images.jpg",
names: Optional[Dict[int, str]] = None,
on_plot: Optional[Callable] = None,
max_size: int = 1920,
max_subplots: int = 16,
save: bool = True,
conf_thres: float = 0.25,
) -> Optional[np.ndarray]:
"""
Plot image grid with labels, bounding boxes, masks, and keypoints.
Args:
labels (Dict[str, Any]): Dictionary containing detection data with keys like 'cls', 'bboxes', 'conf', 'masks', 'keypoints', 'batch_idx', 'img'.
images (torch.Tensor | np.ndarray]): Batch of images to plot. Shape: (batch_size, channels, height, width).
paths (Optional[List[str]]): List of file paths for each image in the batch.
fname (str): Output filename for the plotted image grid.
names (Optional[Dict[int, str]]): Dictionary mapping class indices to class names.
on_plot (Optional[Callable]): Optional callback function to be called after saving the plot.
max_size (int): Maximum size of the output image grid.
max_subplots (int): Maximum number of subplots in the image grid.
save (bool): Whether to save the plotted image grid to a file.
conf_thres (float): Confidence threshold for displaying detections.
Returns:
(np.ndarray): Plotted image grid as a numpy array if save is False, None otherwise.
Note:
This function supports both tensor and numpy array inputs. It will automatically
convert tensor inputs to numpy arrays for processing.
"""
for k in {"cls", "bboxes", "conf", "masks", "keypoints", "batch_idx", "images"}:
if k not in labels:
continue
if k == "cls" and labels[k].ndim == 2:
labels[k] = labels[k].squeeze(1) # squeeze if shape is (n, 1)
if isinstance(labels[k], torch.Tensor):
labels[k] = labels[k].cpu().numpy()
cls = labels.get("cls", np.zeros(0, dtype=np.int64))
batch_idx = labels.get("batch_idx", np.zeros(cls.shape, dtype=np.int64))
bboxes = labels.get("bboxes", np.zeros(0, dtype=np.float32))
confs = labels.get("conf", None)
masks = labels.get("masks", np.zeros(0, dtype=np.uint8))
kpts = labels.get("keypoints", np.zeros(0, dtype=np.float32))
images = labels.get("img", images) # default to input images
if len(images) and isinstance(images, torch.Tensor):
images = images.cpu().float().numpy()
if images.shape[1] > 3:
images = images[:, :3] # crop multispectral images to first 3 channels
bs, _, h, w = images.shape # batch size, _, height, width
bs = min(bs, max_subplots) # limit plot images
ns = np.ceil(bs**0.5) # number of subplots (square)
if np.max(images[0]) <= 1:
images *= 255 # de-normalise (optional)
# Build Image
mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
for i in range(bs):
x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
mosaic[y : y + h, x : x + w, :] = images[i].transpose(1, 2, 0)
# Resize (optional)
scale = max_size / ns / max(h, w)
if scale < 1:
h = math.ceil(scale * h)
w = math.ceil(scale * w)
mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)))
# Annotate
fs = int((h + w) * ns * 0.01) # font size
fs = max(fs, 18) # ensure that the font size is large enough to be easily readable.
annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=str(names))
for i in range(bs):
x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders
if paths:
annotator.text([x + 5, y + 5], text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames
if len(cls) > 0:
idx = batch_idx == i
classes = cls[idx].astype("int")
labels = confs is None
if len(bboxes):
boxes = bboxes[idx]
conf = confs[idx] if confs is not None else None # check for confidence presence (label vs pred)
if len(boxes):
if boxes[:, :4].max() <= 1.1: # if normalized with tolerance 0.1
boxes[..., [0, 2]] *= w # scale to pixels
boxes[..., [1, 3]] *= h
elif scale < 1: # absolute coords need scale if image scales
boxes[..., :4] *= scale
boxes[..., 0] += x
boxes[..., 1] += y
is_obb = boxes.shape[-1] == 5 # xywhr
# TODO: this transformation might be unnecessary
boxes = ops.xywhr2xyxyxyxy(boxes) if is_obb else ops.xywh2xyxy(boxes)
for j, box in enumerate(boxes.astype(np.int64).tolist()):
c = classes[j]
color = colors(c)
c = names.get(c, c) if names else c
if labels or conf[j] > conf_thres:
label = f"{c}" if labels else f"{c} {conf[j]:.1f}"
annotator.box_label(box, label, color=color)
elif len(classes):
for c in classes:
color = colors(c)
c = names.get(c, c) if names else c
annotator.text([x, y], f"{c}", txt_color=color, box_color=(64, 64, 64, 128))
# Plot keypoints
if len(kpts):
kpts_ = kpts[idx].copy()
if len(kpts_):
if kpts_[..., 0].max() <= 1.01 or kpts_[..., 1].max() <= 1.01: # if normalized with tolerance .01
kpts_[..., 0] *= w # scale to pixels
kpts_[..., 1] *= h
elif scale < 1: # absolute coords need scale if image scales
kpts_ *= scale
kpts_[..., 0] += x
kpts_[..., 1] += y
for j in range(len(kpts_)):
if labels or conf[j] > conf_thres:
annotator.kpts(kpts_[j], conf_thres=conf_thres)
# Plot masks
if len(masks):
if idx.shape[0] == masks.shape[0]: # overlap_masks=False
image_masks = masks[idx]
else: # overlap_masks=True
image_masks = masks[[i]] # (1, 640, 640)
nl = idx.sum()
index = np.arange(nl).reshape((nl, 1, 1)) + 1
image_masks = np.repeat(image_masks, nl, axis=0)
image_masks = np.where(image_masks == index, 1.0, 0.0)
im = np.asarray(annotator.im).copy()
for j in range(len(image_masks)):
if labels or conf[j] > conf_thres:
color = colors(classes[j])
mh, mw = image_masks[j].shape
if mh != h or mw != w:
mask = image_masks[j].astype(np.uint8)
mask = cv2.resize(mask, (w, h))
mask = mask.astype(bool)
else:
mask = image_masks[j].astype(bool)
try:
im[y : y + h, x : x + w, :][mask] = (
im[y : y + h, x : x + w, :][mask] * 0.4 + np.array(color) * 0.6
)
except Exception:
pass
annotator.fromarray(im)
if not save:
return np.asarray(annotator.im)
annotator.im.save(fname) # save
if on_plot:
on_plot(fname)
@plt_settings()
def plot_results(
file: str = "path/to/results.csv",
dir: str = "",
segment: bool = False,
pose: bool = False,
classify: bool = False,
on_plot: Optional[Callable] = None,
):
"""
Plot training results from a results CSV file. The function supports various types of data including segmentation,
pose estimation, and classification. Plots are saved as 'results.png' in the directory where the CSV is located.
Args:
file (str, optional): Path to the CSV file containing the training results.
dir (str, optional): Directory where the CSV file is located if 'file' is not provided.
segment (bool, optional): Flag to indicate if the data is for segmentation.
pose (bool, optional): Flag to indicate if the data is for pose estimation.
classify (bool, optional): Flag to indicate if the data is for classification.
on_plot (callable, optional): Callback function to be executed after plotting. Takes filename as an argument.
Examples:
>>> from ultralytics.utils.plotting import plot_results
>>> plot_results("path/to/results.csv", segment=True)
"""
import matplotlib.pyplot as plt # scope for faster 'import ultralytics'
import pandas as pd
from scipy.ndimage import gaussian_filter1d
save_dir = Path(file).parent if file else Path(dir)
if classify:
fig, ax = plt.subplots(2, 2, figsize=(6, 6), tight_layout=True)
index = [2, 5, 3, 4]
elif segment:
fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True)
index = [2, 3, 4, 5, 6, 7, 10, 11, 14, 15, 16, 17, 8, 9, 12, 13]
elif pose:
fig, ax = plt.subplots(2, 9, figsize=(21, 6), tight_layout=True)
index = [2, 3, 4, 5, 6, 7, 8, 11, 12, 15, 16, 17, 18, 19, 9, 10, 13, 14]
else:
fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
index = [2, 3, 4, 5, 6, 9, 10, 11, 7, 8]
ax = ax.ravel()
files = list(save_dir.glob("results*.csv"))
assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot."
for f in files:
try:
data = pd.read_csv(f)
s = [x.strip() for x in data.columns]
x = data.values[:, 0]
for i, j in enumerate(index):
y = data.values[:, j].astype("float")
# y[y == 0] = np.nan # don't show zero values
ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=8) # actual results
ax[i].plot(x, gaussian_filter1d(y, sigma=3), ":", label="smooth", linewidth=2) # smoothing line
ax[i].set_title(s[j], fontsize=12)
# if j in {8, 9, 10}: # share train and val loss y axes
# ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
except Exception as e:
LOGGER.error(f"Plotting error for {f}: {e}")
ax[1].legend()
fname = save_dir / "results.png"
fig.savefig(fname, dpi=200)
plt.close()
if on_plot:
on_plot(fname)
def plt_color_scatter(v, f, bins: int = 20, cmap: str = "viridis", alpha: float = 0.8, edgecolors: str = "none"):
"""
Plot a scatter plot with points colored based on a 2D histogram.
Args:
v (array-like): Values for the x-axis.
f (array-like): Values for the y-axis.
bins (int, optional): Number of bins for the histogram.
cmap (str, optional): Colormap for the scatter plot.
alpha (float, optional): Alpha for the scatter plot.
edgecolors (str, optional): Edge colors for the scatter plot.
Examples:
>>> v = np.random.rand(100)
>>> f = np.random.rand(100)
>>> plt_color_scatter(v, f)
"""
import matplotlib.pyplot as plt # scope for faster 'import ultralytics'
# Calculate 2D histogram and corresponding colors
hist, xedges, yedges = np.histogram2d(v, f, bins=bins)
colors = [
hist[
min(np.digitize(v[i], xedges, right=True) - 1, hist.shape[0] - 1),
min(np.digitize(f[i], yedges, right=True) - 1, hist.shape[1] - 1),
]
for i in range(len(v))
]
# Scatter plot
plt.scatter(v, f, c=colors, cmap=cmap, alpha=alpha, edgecolors=edgecolors)
def plot_tune_results(csv_file: str = "tune_results.csv"):
"""
Plot the evolution results stored in a 'tune_results.csv' file. The function generates a scatter plot for each key
in the CSV, color-coded based on fitness scores. The best-performing configurations are highlighted on the plots.
Args:
csv_file (str, optional): Path to the CSV file containing the tuning results.
Examples:
>>> plot_tune_results("path/to/tune_results.csv")
"""
import matplotlib.pyplot as plt # scope for faster 'import ultralytics'
import pandas as pd
from scipy.ndimage import gaussian_filter1d
def _save_one_file(file):
"""Save one matplotlib plot to 'file'."""
plt.savefig(file, dpi=200)
plt.close()
LOGGER.info(f"Saved {file}")
# Scatter plots for each hyperparameter
csv_file = Path(csv_file)
data = pd.read_csv(csv_file)
num_metrics_columns = 1
keys = [x.strip() for x in data.columns][num_metrics_columns:]
x = data.values
fitness = x[:, 0] # fitness
j = np.argmax(fitness) # max fitness index
n = math.ceil(len(keys) ** 0.5) # columns and rows in plot
plt.figure(figsize=(10, 10), tight_layout=True)
for i, k in enumerate(keys):
v = x[:, i + num_metrics_columns]
mu = v[j] # best single result
plt.subplot(n, n, i + 1)
plt_color_scatter(v, fitness, cmap="viridis", alpha=0.8, edgecolors="none")
plt.plot(mu, fitness.max(), "k+", markersize=15)
plt.title(f"{k} = {mu:.3g}", fontdict={"size": 9}) # limit to 40 characters
plt.tick_params(axis="both", labelsize=8) # Set axis label size to 8
if i % n != 0:
plt.yticks([])
_save_one_file(csv_file.with_name("tune_scatter_plots.png"))
# Fitness vs iteration
x = range(1, len(fitness) + 1)
plt.figure(figsize=(10, 6), tight_layout=True)
plt.plot(x, fitness, marker="o", linestyle="none", label="fitness")
plt.plot(x, gaussian_filter1d(fitness, sigma=3), ":", label="smoothed", linewidth=2) # smoothing line
plt.title("Fitness vs Iteration")
plt.xlabel("Iteration")
plt.ylabel("Fitness")
plt.grid(True)
plt.legend()
_save_one_file(csv_file.with_name("tune_fitness.png"))
def feature_visualization(x, module_type: str, stage: int, n: int = 32, save_dir: Path = Path("runs/detect/exp")):
"""
Visualize feature maps of a given model module during inference.
Args:
x (torch.Tensor): Features to be visualized.
module_type (str): Module type.
stage (int): Module stage within the model.
n (int, optional): Maximum number of feature maps to plot.
save_dir (Path, optional): Directory to save results.
"""
import matplotlib.pyplot as plt # scope for faster 'import ultralytics'
for m in {"Detect", "Segment", "Pose", "Classify", "OBB", "RTDETRDecoder"}: # all model heads
if m in module_type:
return
if isinstance(x, torch.Tensor):
_, channels, height, width = x.shape # batch, channels, height, width
if height > 1 and width > 1:
f = save_dir / f"stage{stage}_{module_type.rsplit('.', 1)[-1]}_features.png" # filename
blocks = torch.chunk(x[0].cpu(), channels, dim=0) # select batch index 0, block by channels
n = min(n, channels) # number of plots
_, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True) # 8 rows x n/8 cols
ax = ax.ravel()
plt.subplots_adjust(wspace=0.05, hspace=0.05)
for i in range(n):
ax[i].imshow(blocks[i].squeeze()) # cmap='gray'
ax[i].axis("off")
LOGGER.info(f"Saving {f}... ({n}/{channels})")
plt.savefig(f, dpi=300, bbox_inches="tight")
plt.close()
np.save(str(f.with_suffix(".npy")), x[0].cpu().numpy()) # npy save