Reference for ultralytics/utils/plotting.py

def __call__(self, i, bgr=False):
    """Converts hex color codes to RGB values."""
    c = self.palette[int(i) % self.n]
    return (c[2], c[1], c[0]) if bgr else c

@staticmethod
def hex2rgb(h):
    """Converts 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))

def box_label(self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255), rotated=False):
    """
    Draws a bounding box to image with label.

    Args:
        box (tuple): The bounding box coordinates (x1, y1, x2, y2).
        label (str): 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).
        rotated (bool, optional): Variable used to check if task is OBB
    """
    txt_color = self.get_txt_color(color, txt_color)
    if isinstance(box, torch.Tensor):
        box = box.tolist()
    if self.pil or not is_ascii(label):
        if rotated:
            p1 = box[0]
            self.draw.polygon([tuple(b) for b in box], width=self.lw, outline=color)  # PIL requires tuple box
        else:
            p1 = (box[0], box[1])
            self.draw.rectangle(box, width=self.lw, outline=color)  # box
        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
        if rotated:
            p1 = [int(b) for b in box[0]]
            cv2.polylines(self.im, [np.asarray(box, dtype=int)], True, color, self.lw)  # cv2 requires nparray box
        else:
            p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
            cv2.rectangle(self.im, p1, p2, 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 circle_label(self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255), margin=2):
    """
    Draws a label with a background circle centered within a given bounding box.

    Args:
        box (tuple): The bounding box coordinates (x1, y1, x2, y2).
        label (str): 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).
        margin (int, optional): The margin between the text and the rectangle border.
    """
    # If label have more than 3 characters, skip other characters, due to circle size
    if len(label) > 3:
        print(
            f"Length of label is {len(label)}, initial 3 label characters will be considered for circle annotation!"
        )
        label = label[:3]

    # Calculate the center of the box
    x_center, y_center = int((box[0] + box[2]) / 2), int((box[1] + box[3]) / 2)
    # Get the text size
    text_size = cv2.getTextSize(str(label), cv2.FONT_HERSHEY_SIMPLEX, self.sf - 0.15, self.tf)[0]
    # Calculate the required radius to fit the text with the margin
    required_radius = int(((text_size[0] ** 2 + text_size[1] ** 2) ** 0.5) / 2) + margin
    # Draw the circle with the required radius
    cv2.circle(self.im, (x_center, y_center), required_radius, color, -1)
    # Calculate the position for the text
    text_x = x_center - text_size[0] // 2
    text_y = y_center + text_size[1] // 2
    # Draw the text
    cv2.putText(
        self.im,
        str(label),
        (text_x, text_y),
        cv2.FONT_HERSHEY_SIMPLEX,
        self.sf - 0.15,
        self.get_txt_color(color, txt_color),
        self.tf,
        lineType=cv2.LINE_AA,
    )

def display_analytics(self, im0, text, txt_color, bg_color, margin):
    """
    Display the overall statistics for parking lots.

    Args:
        im0 (ndarray): Inference image.
        text (dict): Labels dictionary.
        txt_color (tuple): Display color for text foreground.
        bg_color (tuple): Display color for text background.
        margin (int): Gap between text and rectangle for better display.
    """
    horizontal_gap = int(im0.shape[1] * 0.02)
    vertical_gap = int(im0.shape[0] * 0.01)
    text_y_offset = 0
    for label, value in text.items():
        txt = f"{label}: {value}"
        text_size = cv2.getTextSize(txt, 0, self.sf, self.tf)[0]
        if text_size[0] < 5 or text_size[1] < 5:
            text_size = (5, 5)
        text_x = im0.shape[1] - text_size[0] - margin * 2 - horizontal_gap
        text_y = text_y_offset + text_size[1] + margin * 2 + vertical_gap
        rect_x1 = text_x - margin * 2
        rect_y1 = text_y - text_size[1] - margin * 2
        rect_x2 = text_x + text_size[0] + margin * 2
        rect_y2 = text_y + margin * 2
        cv2.rectangle(im0, (rect_x1, rect_y1), (rect_x2, rect_y2), bg_color, -1)
        cv2.putText(im0, txt, (text_x, text_y), 0, self.sf, txt_color, self.tf, lineType=cv2.LINE_AA)
        text_y_offset = rect_y2

def display_objects_labels(self, im0, text, txt_color, bg_color, x_center, y_center, margin):
    """
    Display the bounding boxes labels in parking management app.

    Args:
        im0 (ndarray): Inference image.
        text (str): Object/class name.
        txt_color (tuple): Display color for text foreground.
        bg_color (tuple): Display color for text background.
        x_center (float): The x position center point for bounding box.
        y_center (float): The y position center point for bounding box.
        margin (int): The gap between text and rectangle for better display.
    """
    text_size = cv2.getTextSize(text, 0, fontScale=self.sf, thickness=self.tf)[0]
    text_x = x_center - text_size[0] // 2
    text_y = y_center + text_size[1] // 2

    rect_x1 = text_x - margin
    rect_y1 = text_y - text_size[1] - margin
    rect_x2 = text_x + text_size[0] + margin
    rect_y2 = text_y + margin
    cv2.rectangle(im0, (rect_x1, rect_y1), (rect_x2, rect_y2), bg_color, -1)
    cv2.putText(im0, text, (text_x, text_y), 0, self.sf, txt_color, self.tf, lineType=cv2.LINE_AA)

def draw_centroid_and_tracks(self, track, color=(255, 0, 255), track_thickness=2):
    """
    Draw centroid point and track trails.

    Args:
        track (list): object tracking points for trails display
        color (tuple): tracks line color
        track_thickness (int): track line thickness value
    """
    points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
    cv2.polylines(self.im, [points], isClosed=False, color=color, thickness=track_thickness)
    cv2.circle(self.im, (int(track[-1][0]), int(track[-1][1])), track_thickness * 2, color, -1)

def draw_region(self, reg_pts=None, color=(0, 255, 0), thickness=5):
    """
    Draw region line.

    Args:
        reg_pts (list): Region Points (for line 2 points, for region 4 points)
        color (tuple): Region Color value
        thickness (int): Region area thickness value
    """
    cv2.polylines(self.im, [np.array(reg_pts, dtype=np.int32)], isClosed=True, color=color, thickness=thickness)

    # Draw small circles at the corner points
    for point in reg_pts:
        cv2.circle(self.im, (point[0], point[1]), thickness * 2, color, -1)  # -1 fills the circle

def draw_specific_points(self, keypoints, indices=None, radius=2, conf_thres=0.25):
    """
    Draw specific keypoints for gym steps counting.

    Args:
        keypoints (list): Keypoints data to be plotted.
        indices (list, optional): Keypoint indices to be plotted. Defaults to [2, 5, 7].
        radius (int, optional): Keypoint radius. Defaults to 2.
        conf_thres (float, optional): Confidence threshold for keypoints. Defaults to 0.25.

    Returns:
        (numpy.ndarray): Image with drawn keypoints.

    Note:
        Keypoint format: [x, y] or [x, y, confidence].
        Modifies self.im in-place.
    """
    indices = indices or [2, 5, 7]
    points = [(int(k[0]), int(k[1])) for i, k in enumerate(keypoints) if i in indices and k[2] >= conf_thres]

    # Draw lines between consecutive points
    for start, end in zip(points[:-1], points[1:]):
        cv2.line(self.im, start, end, (0, 255, 0), 2, lineType=cv2.LINE_AA)

    # Draw circles for keypoints
    for pt in points:
        cv2.circle(self.im, pt, radius, (0, 0, 255), -1, lineType=cv2.LINE_AA)

    return self.im

@staticmethod
def estimate_pose_angle(a, b, c):
    """
    Calculate the pose angle for object.

    Args:
        a (float) : The value of pose point a
        b (float): The value of pose point b
        c (float): The value o pose point c

    Returns:
        angle (degree): Degree value of angle between three points
    """
    a, b, c = np.array(a), np.array(b), np.array(c)
    radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
    angle = np.abs(radians * 180.0 / np.pi)
    if angle > 180.0:
        angle = 360 - angle
    return angle

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 get_bbox_dimension(self, bbox=None):
    """
    Calculate the 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.
    """
    x_min, y_min, x_max, y_max = bbox
    width = x_max - x_min
    height = y_max - y_min
    return width, height, width * height

def get_txt_color(self, color=(128, 128, 128), txt_color=(255, 255, 255)):
    """
    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:
        txt_color (tuple): Text color for label
    """
    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 kpts(self, kpts, shape=(640, 640), radius=None, kpt_line=True, conf_thres=0.25, kpt_color=None):
    """
    Plot keypoints on the image.

    Args:
        kpts (torch.Tensor): Keypoints, shape [17, 3] (x, y, confidence).
        shape (tuple, optional): Image shape (h, w). Defaults to (640, 640).
        radius (int, optional): Keypoint radius. Defaults to 5.
        kpt_line (bool, optional): Draw lines between keypoints. Defaults to True.
        conf_thres (float, optional): Confidence threshold. Defaults to 0.25.
        kpt_color (tuple, optional): Keypoint color (B, G, R). Defaults to None.

    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 masks(self, masks, colors, im_gpu, alpha=0.5, retina_masks=False):
    """
    Plot masks on image.

    Args:
        masks (tensor): Predicted masks on cuda, shape: [n, h, w]
        colors (List[List[Int]]): Colors for predicted masks, [[r, g, b] * n]
        im_gpu (tensor): Image is in cuda, shape: [3, h, w], range: [0, 1]
        alpha (float): Mask transparency: 0.0 fully transparent, 1.0 opaque
        retina_masks (bool): Whether to use high resolution masks or not. Defaults to False.
    """
    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 plot_angle_and_count_and_stage(
    self, angle_text, count_text, stage_text, center_kpt, color=(104, 31, 17), txt_color=(255, 255, 255)
):
    """
    Plot the pose angle, count value, and step stage.

    Args:
        angle_text (str): Angle value for workout monitoring
        count_text (str): Counts value for workout monitoring
        stage_text (str): Stage decision for workout monitoring
        center_kpt (list): Centroid pose index for workout monitoring
        color (tuple, optional): Text background color
        txt_color (tuple, optional): Text foreground color
    """
    # Format text
    angle_text, count_text, stage_text = f" {angle_text:.2f}", f"Steps : {count_text}", f" {stage_text}"

    # Draw angle, count and stage text
    angle_height = self.plot_workout_information(
        angle_text, (int(center_kpt[0]), int(center_kpt[1])), color, txt_color
    )
    count_height = self.plot_workout_information(
        count_text, (int(center_kpt[0]), int(center_kpt[1]) + angle_height + 20), color, txt_color
    )
    self.plot_workout_information(
        stage_text, (int(center_kpt[0]), int(center_kpt[1]) + angle_height + count_height + 40), color, txt_color
    )

def plot_distance_and_line(
    self, pixels_distance, centroids, line_color=(104, 31, 17), centroid_color=(255, 0, 255)
):
    """
    Plot the distance and line on frame.

    Args:
        pixels_distance (float): Pixels distance between two bbox centroids.
        centroids (list): Bounding box centroids data.
        line_color (tuple, optional): Distance line color.
        centroid_color (tuple, optional): Bounding box centroid color.
    """
    # Get the text size
    text = f"Pixels Distance: {pixels_distance:.2f}"
    (text_width_m, text_height_m), _ = cv2.getTextSize(text, 0, self.sf, self.tf)

    # Define corners with 10-pixel margin and draw rectangle
    cv2.rectangle(self.im, (15, 25), (15 + text_width_m + 20, 25 + text_height_m + 20), line_color, -1)

    # Calculate the position for the text with a 10-pixel margin and draw text
    text_position = (25, 25 + text_height_m + 10)
    cv2.putText(
        self.im,
        text,
        text_position,
        0,
        self.sf,
        (255, 255, 255),
        self.tf,
        cv2.LINE_AA,
    )

    cv2.line(self.im, centroids[0], centroids[1], line_color, 3)
    cv2.circle(self.im, centroids[0], 6, centroid_color, -1)
    cv2.circle(self.im, centroids[1], 6, centroid_color, -1)

def plot_workout_information(self, display_text, position, color=(104, 31, 17), txt_color=(255, 255, 255)):
    """
    Draw text with a background on the image.

    Args:
        display_text (str): The text to be displayed.
        position (tuple): Coordinates (x, y) on the image where the text will be placed.
        color (tuple, optional): Text background color
        txt_color (tuple, optional): Text foreground color
    """
    (text_width, text_height), _ = cv2.getTextSize(display_text, 0, self.sf, self.tf)

    # Draw background rectangle
    cv2.rectangle(
        self.im,
        (position[0], position[1] - text_height - 5),
        (position[0] + text_width + 10, position[1] - text_height - 5 + text_height + 10 + self.tf),
        color,
        -1,
    )
    # Draw text
    cv2.putText(self.im, display_text, position, 0, self.sf, txt_color, self.tf)

    return text_height