From ba559c2658d467697cb6fe3322280ec6c77731dd Mon Sep 17 00:00:00 2001 From: KshitijAucharmal Date: Mon, 28 Oct 2024 08:59:30 +0530 Subject: [PATCH 01/16] Added smart positioning (non overlapping labels) to VertexLabelAnnotator --- supervision/detection/utils.py | 83 ++++++++++++++++++++++++++++++ supervision/keypoint/annotators.py | 12 ++++- 2 files changed, 94 insertions(+), 1 deletion(-) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 43fcec5a0..308797eab 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1,3 +1,4 @@ +import math from itertools import chain from typing import Dict, List, Optional, Tuple, Union @@ -1039,3 +1040,85 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: ) vector_start = np.array([vector.start.x, vector.start.y]) return np.cross(vector_at_zero, anchors - vector_start) + + +# Intelligent padding functions +def get_intersection_center( + xyxy_1: np.ndarray, xyxy_2: np.ndarray +) -> Optional[Tuple[float, float]]: + overlap_xmin = max(xyxy_1[0], xyxy_2[0]) + overlap_ymin = max(xyxy_1[1], xyxy_2[1]) + overlap_xmax = min(xyxy_1[2], xyxy_2[2]) + overlap_ymax = min(xyxy_1[3], xyxy_2[3]) + + if overlap_xmin < overlap_xmax and overlap_ymin < overlap_ymax: + x_center = (overlap_xmin + overlap_xmax) / 2 + y_center = (overlap_ymin + overlap_ymax) / 2 + return (x_center, y_center) + else: + return None + + +def get_box_center(xyxy: np.ndarray) -> Tuple[float, float]: + x_center = (xyxy[0] + xyxy[2]) / 2 + y_center = (xyxy[1] + xyxy[3]) / 2 + return (x_center, y_center) + + +def vector_with_length( + xy_1: Tuple[float, float], xy_2: Tuple[float, float], n: float +) -> Tuple[float, float]: + x1, y1 = xy_1 + x2, y2 = xy_2 + + dx = x2 - x1 + dy = y2 - y1 + + if dx == 0 and dy == 0: + return 0, 0 + + magnitude = math.sqrt(dx**2 + dy**2) + + unit_dx = dx / magnitude + unit_dy = dy / magnitude + + v1 = unit_dx * n + v2 = unit_dy * n + + return (v1, v2) + + +def pad(xyxy: np.ndarray, px: int, py: Optional[int] = None): + if py is None: + py = px + + result = xyxy.copy() + result[:, [0, 1]] -= [px, py] + result[:, [2, 3]] += [px, py] + + return result + + +def spread_out(xyxy: np.ndarray, step) -> np.ndarray: + xyxy_padded = pad(xyxy, px=step) + while True: + iou = box_iou_batch(xyxy_padded, xyxy_padded) + np.fill_diagonal(iou, 0) + + if np.all(iou == 0): + return pad(xyxy_padded, px=-step) + + i, j = np.unravel_index(np.argmax(iou), iou.shape) + + xyxy_i, xyxy_j = xyxy_padded[i], xyxy_padded[j] + intersection_center = get_intersection_center(xyxy_i, xyxy_j) + xyxy_i_center = get_box_center(xyxy_i) + xyxy_j_center = get_box_center(xyxy_j) + + vector_i = vector_with_length(intersection_center, xyxy_i_center, step) + vector_j = vector_with_length(intersection_center, xyxy_j_center, step) + + xyxy_padded[i, [0, 2]] += int(vector_i[0]) + xyxy_padded[i, [1, 3]] += int(vector_i[1]) + xyxy_padded[j, [0, 2]] += int(vector_j[0]) + xyxy_padded[j, [1, 3]] += int(vector_j[1]) diff --git a/supervision/keypoint/annotators.py b/supervision/keypoint/annotators.py index 559bfa921..e82acbc65 100644 --- a/supervision/keypoint/annotators.py +++ b/supervision/keypoint/annotators.py @@ -7,6 +7,7 @@ from supervision import Rect, pad_boxes from supervision.annotators.base import ImageType +from supervision.detection.utils import pad, spread_out from supervision.draw.color import Color from supervision.draw.utils import draw_rounded_rectangle from supervision.keypoint.core import KeyPoints @@ -201,6 +202,7 @@ def __init__( text_thickness: int = 1, text_padding: int = 10, border_radius: int = 0, + use_smart_positioning: bool = False, ): """ Args: @@ -215,7 +217,10 @@ def __init__( text_padding (int): The padding around the text. border_radius (int): The radius of the rounded corners of the boxes. Set to a high value to produce circles. + use_smart_positioning (bool): Whether to use smart positioning to prevent + label overlapping or not. """ + self.use_smart_positioning = use_smart_positioning self.border_radius: int = border_radius self.color: Union[Color, List[Color]] = color self.text_color: Union[Color, List[Color]] = text_color @@ -357,7 +362,12 @@ def annotate( ] ) - xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding) + if self.use_smart_positioning: + xyxy_padded = pad(xyxy=xyxy, px=self.text_padding) + xyxy_padded = spread_out(xyxy_padded, step=2) + xyxy = pad(xyxy=xyxy_padded, px=-self.text_padding) + else: + xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding) for text, color, text_color, box, box_padded in zip( labels, colors, text_colors, xyxy, xyxy_padded From a2e2e3725a17de63d5d2c7e97f79722f9880753d Mon Sep 17 00:00:00 2001 From: KshitijAucharmal Date: Mon, 28 Oct 2024 19:05:39 +0530 Subject: [PATCH 02/16] Added smart positioning (non overlapping labels) to LabelAnnotator --- supervision/annotators/core.py | 176 ++++++++++++++++++++++++--------- 1 file changed, 129 insertions(+), 47 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 1910ac9f4..9d08cf1df 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -16,7 +16,7 @@ ) from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES from supervision.detection.core import Detections -from supervision.detection.utils import clip_boxes, mask_to_polygons +from supervision.detection.utils import clip_boxes, mask_to_polygons, spread_out from supervision.draw.color import Color, ColorPalette from supervision.draw.utils import draw_polygon from supervision.geometry.core import Position @@ -1054,6 +1054,7 @@ def __init__( text_position: Position = Position.TOP_LEFT, color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, + use_smart_positioning: bool = False, ): """ Args: @@ -1070,7 +1071,10 @@ def __init__( Options are `INDEX`, `CLASS`, `TRACK`. border_radius (int): The radius to apply round edges. If the selected value is higher than the lower dimension, width or height, is clipped. + use_smart_positioning (bool): Whether to use smart positioning to prevent + label overlapping or not. """ + self.use_smart_positioning: bool = use_smart_positioning self.border_radius: int = border_radius self.color: Union[Color, ColorPalette] = color self.text_color: Union[Color, ColorPalette] = text_color @@ -1128,11 +1132,35 @@ def annotate( ![label-annotator-example](https://media.roboflow.com/ supervision-annotator-examples/label-annotator-example-purple.png) """ + assert isinstance(scene, np.ndarray) - font = cv2.FONT_HERSHEY_SIMPLEX - anchors_coordinates = detections.get_anchors_coordinates( - anchor=self.text_anchor - ).astype(int) + self._validate_labels(labels, detections) + + # Get text properties for all detections + text_props = self._get_text_properties(detections, labels) + + # Calculate background coordinates for all labels + xyxy = self._calculate_label_backgrounds( + detections, text_props, self.text_anchor, self.text_padding + ) + + # Adjust positions if smart positioning is enabled + if self.use_smart_positioning: + xyxy = spread_out(xyxy, step=2) + + # Draw all labels + self._draw_labels( + scene=scene, + xyxy=xyxy, + text_props=text_props, + detections=detections, + custom_color_lookup=custom_color_lookup, + ) + + return scene + + def _validate_labels(self, labels: Optional[List[str]], detections: Detections): + """Validates that the number of labels matches the number of detections.""" if labels is not None and len(labels) != len(detections): raise ValueError( f"The number of labels ({len(labels)}) does not match the " @@ -1140,64 +1168,119 @@ def annotate( f"should have exactly 1 label." ) - for detection_idx, center_coordinates in enumerate(anchors_coordinates): - color = resolve_color( - color=self.color, - detections=detections, - detection_idx=detection_idx, - color_lookup=( - self.color_lookup - if custom_color_lookup is None - else custom_color_lookup - ), - ) - - text_color = resolve_color( - color=self.text_color, - detections=detections, - detection_idx=detection_idx, - color_lookup=( - self.color_lookup - if custom_color_lookup is None - else custom_color_lookup - ), - ) + def _get_text_properties( + self, detections: Detections, custom_labels: Optional[List[str]] + ) -> List[dict]: + """Gets text content and dimensions for all detections.""" + text_props = [] + font = cv2.FONT_HERSHEY_SIMPLEX - if labels is not None: - text = labels[detection_idx] - elif CLASS_NAME_DATA_FIELD in detections.data: - text = detections.data[CLASS_NAME_DATA_FIELD][detection_idx] - elif detections.class_id is not None: - text = str(detections.class_id[detection_idx]) - else: - text = str(detection_idx) + for idx in range(len(detections)): + # Determine label text + text = self._get_label_text(detections, custom_labels, idx) - text_w, text_h = cv2.getTextSize( + # Calculate text dimensions + (text_w, text_h) = cv2.getTextSize( text=text, fontFace=font, fontScale=self.text_scale, thickness=self.text_thickness, )[0] - text_w_padded = text_w + 2 * self.text_padding - text_h_padded = text_h + 2 * self.text_padding + + text_props.append( + { + "text": text, + "width": text_w, + "height": text_h, + "width_padded": text_w + 2 * self.text_padding, + "height_padded": text_h + 2 * self.text_padding, + } + ) + + return text_props + + def _get_label_text( + self, detections: Detections, custom_labels: Optional[List[str]], idx: int + ) -> str: + """Determines the label text for a given detection.""" + if custom_labels is not None: + return custom_labels[idx] + elif CLASS_NAME_DATA_FIELD in detections.data: + return detections.data[CLASS_NAME_DATA_FIELD][idx] + elif detections.class_id is not None: + return str(detections.class_id[idx]) + return str(idx) + + def _calculate_label_backgrounds( + self, + detections: Detections, + text_props: List[dict], + text_anchor: str, + text_padding: int, + ) -> np.ndarray: + """Calculates background coordinates for all labels.""" + anchors_coordinates = detections.get_anchors_coordinates( + anchor=text_anchor + ).astype(int) + + xyxy = [] + for idx, center_coords in enumerate(anchors_coordinates): text_background_xyxy = resolve_text_background_xyxy( - center_coordinates=tuple(center_coordinates), - text_wh=(text_w_padded, text_h_padded), - position=self.text_anchor, + center_coordinates=tuple(center_coords), + text_wh=( + text_props[idx]["width_padded"], + text_props[idx]["height_padded"], + ), + position=text_anchor, + ) + xyxy.append(text_background_xyxy) + + return np.array(xyxy) + + def _draw_labels( + self, + scene: ImageType, + xyxy: np.ndarray, + text_props: List[dict], + detections: Detections, + custom_color_lookup: Optional[np.ndarray], + ) -> None: + """Draws all labels and their backgrounds on the scene.""" + if custom_color_lookup is not None: + color_lookup = custom_color_lookup + else: + color_lookup = self.color_lookup + font = cv2.FONT_HERSHEY_SIMPLEX + + for idx, coordinates in enumerate(xyxy): + # Resolve colors + bg_color = resolve_color( + color=self.color, + detections=detections, + detection_idx=idx, + color_lookup=color_lookup, + ) + text_color = resolve_color( + color=self.text_color, + detections=detections, + detection_idx=idx, + color_lookup=color_lookup, ) - text_x = text_background_xyxy[0] + self.text_padding - text_y = text_background_xyxy[1] + self.text_padding + text_h + # Calculate text position + text_x = coordinates[0] + self.text_padding + text_y = coordinates[1] + self.text_padding + text_props[idx]["height"] + # Draw background and text self.draw_rounded_rectangle( scene=scene, - xyxy=text_background_xyxy, - color=color.as_bgr(), + xyxy=coordinates, + color=bg_color.as_bgr(), border_radius=self.border_radius, ) cv2.putText( img=scene, - text=text, + text=text_props[idx]["text"], org=(text_x, text_y), fontFace=font, fontScale=self.text_scale, @@ -1205,7 +1288,6 @@ def annotate( thickness=self.text_thickness, lineType=cv2.LINE_AA, ) - return scene @staticmethod def draw_rounded_rectangle( From ecf5b1334fb1042565a945ed48285147cfc9dbd1 Mon Sep 17 00:00:00 2001 From: KshitijAucharmal Date: Tue, 29 Oct 2024 16:06:54 +0530 Subject: [PATCH 03/16] Added smart positioning (non overlapping labels) to RichLabelAnnotator --- supervision/annotators/core.py | 157 +++++++++++++++++++++++---------- 1 file changed, 112 insertions(+), 45 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 9d08cf1df..167782bb6 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1348,6 +1348,7 @@ def __init__( text_position: Position = Position.TOP_LEFT, color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, + use_smart_positioning: bool = False, ): """ Args: @@ -1364,6 +1365,8 @@ def __init__( Options are `INDEX`, `CLASS`, `TRACK`. border_radius (int): The radius to apply round edges. If the selected value is higher than the lower dimension, width or height, is clipped. + use_smart_positioning (bool): Whether to use smart positioning to prevent + label overlapping or not. """ self.color = color self.text_color = text_color @@ -1371,6 +1374,7 @@ def __init__( self.text_anchor = text_position self.color_lookup = color_lookup self.border_radius = border_radius + self.use_smart_positioning: bool = use_smart_positioning if font_path is not None: try: self.font = ImageFont.truetype(font_path, font_size) @@ -1429,72 +1433,135 @@ def annotate( """ assert isinstance(scene, Image.Image) draw = ImageDraw.Draw(scene) - anchors_coordinates = detections.get_anchors_coordinates( - anchor=self.text_anchor - ).astype(int) + + # Input validation if labels is not None and len(labels) != len(detections): raise ValueError( - f"The number of labels provided ({len(labels)}) does not match the " - f"number of detections ({len(detections)}). Each detection should have " - f"a corresponding label." + f"Label count ({len(labels)}) != detection count ({len(detections)})" ) - for detection_idx, center_coordinates in enumerate(anchors_coordinates): - color = resolve_color( + + # Get anchor coordinates for all detections + detection_anchor_coordinates = detections.get_anchors_coordinates( + anchor=self.text_anchor + ).astype(int) + + # Use the appropriate color lookup table + effective_color_lookup = ( + custom_color_lookup + if custom_color_lookup is not None + else self.color_lookup + ) + + def _get_detection_label_text(detection_index: int) -> str: + """ + Determine the appropriate label text for a detection. + Args: + detection_index: Index of the current detection + Returns: + str: The label text to display + """ + if labels is not None: + return labels[detection_index] + if CLASS_NAME_DATA_FIELD in detections.data: + return detections.data[CLASS_NAME_DATA_FIELD][detection_index] + if detections.class_id is not None: + return str(detections.class_id[detection_index]) + return str(detection_index) + + def _calculate_text_dimensions(label_text: str) -> tuple: + """ + Calculate text dimensions and offsets for the given label text. + Args: + label_text: The text to measure + Returns: + tuple: ((left_offset, top_offset), (padded_width, padded_height)) + """ + text_left, text_top, text_right, text_bottom = draw.textbbox( + (0, 0), label_text, font=self.font + ) + text_width = text_right - text_left + text_height = text_bottom - text_top + padded_width = text_width + 2 * self.text_padding + padded_height = text_height + 2 * self.text_padding + return (text_left, text_top), (padded_width, padded_height) + + # Prepare all annotation data + annotation_collection = [] + for detection_index, center_coordinate in enumerate( + detection_anchor_coordinates + ): + # Get colors once per detection + background_color = resolve_color( color=self.color, detections=detections, - detection_idx=detection_idx, - color_lookup=( - self.color_lookup - if custom_color_lookup is None - else custom_color_lookup - ), + detection_idx=detection_index, + color_lookup=effective_color_lookup, ) - - text_color = resolve_color( + label_text_color = resolve_color( color=self.text_color, detections=detections, - detection_idx=detection_idx, - color_lookup=( - self.color_lookup - if custom_color_lookup is None - else custom_color_lookup - ), + detection_idx=detection_index, + color_lookup=effective_color_lookup, ) - if labels is not None: - text = labels[detection_idx] - elif CLASS_NAME_DATA_FIELD in detections.data: - text = detections.data[CLASS_NAME_DATA_FIELD][detection_idx] - elif detections.class_id is not None: - text = str(detections.class_id[detection_idx]) - else: - text = str(detection_idx) + # Get text and calculate dimensions + label_text = _get_detection_label_text(detection_index) + text_offset_coordinates, padded_dimensions = _calculate_text_dimensions( + label_text + ) - left, top, right, bottom = draw.textbbox((0, 0), text, font=self.font) - text_width = right - left - text_height = bottom - top - text_w_padded = text_width + 2 * self.text_padding - text_h_padded = text_height + 2 * self.text_padding - text_background_xyxy = resolve_text_background_xyxy( - center_coordinates=tuple(center_coordinates), - text_wh=(text_w_padded, text_h_padded), + # Calculate background coordinates + background_coordinates = resolve_text_background_xyxy( + center_coordinates=tuple(center_coordinate), + text_wh=padded_dimensions, position=self.text_anchor, ) - text_x = text_background_xyxy[0] + self.text_padding - left - text_y = text_background_xyxy[1] + self.text_padding - top + # Store all data for this annotation + annotation_collection.append( + { + "label_text": label_text, + "background_color": background_color, + "text_color": label_text_color, + "text_offset": text_offset_coordinates, + "background_coordinates": background_coordinates, + } + ) + + # Convert coordinates to numpy array for processing + background_coordinate_array = np.array( + [data["background_coordinates"] for data in annotation_collection] + ) + + # Apply smart positioning if enabled + if self.use_smart_positioning: + background_coordinate_array = spread_out( + background_coordinate_array, step=2 + ) + + # Draw annotations + for annotation_index, coordinates in enumerate(background_coordinate_array): + annotation_data = annotation_collection[annotation_index] + + # Calculate final text position + label_x_position = ( + coordinates[0] + self.text_padding - annotation_data["text_offset"][0] + ) + label_y_position = ( + coordinates[1] + self.text_padding - annotation_data["text_offset"][1] + ) draw.rounded_rectangle( - text_background_xyxy, + tuple(coordinates), radius=self.border_radius, - fill=color.as_rgb(), + fill=annotation_data["background_color"].as_rgb(), outline=None, ) draw.text( - xy=(text_x, text_y), - text=text, + xy=(label_x_position, label_y_position), + text=annotation_data["label_text"], font=self.font, - fill=text_color.as_rgb(), + fill=annotation_data["text_color"].as_rgb(), ) return scene From 6728527f36b6ecc728d35232c58240f2ffd3cb98 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Thu, 7 Nov 2024 17:23:23 +0200 Subject: [PATCH 04/16] Smart annotators are now called via "spread_out" param. * Alos, refactored RichLabelAnnotator so its structure matches LabelAnnotator --- supervision/annotators/core.py | 432 ++++++++++++++++------------- supervision/detection/utils.py | 86 ++---- supervision/keypoint/annotators.py | 49 ++-- 3 files changed, 287 insertions(+), 280 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 167782bb6..2864bfdd9 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,6 +1,7 @@ +from dataclasses import dataclass from functools import lru_cache from math import sqrt -from typing import List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union import cv2 import numpy as np @@ -16,7 +17,7 @@ ) from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES from supervision.detection.core import Detections -from supervision.detection.utils import clip_boxes, mask_to_polygons, spread_out +from supervision.detection.utils import clip_boxes, mask_to_polygons, spread_out_boxes from supervision.draw.color import Color, ColorPalette from supervision.draw.utils import draw_polygon from supervision.geometry.core import Position @@ -1044,6 +1045,16 @@ class LabelAnnotator(BaseAnnotator): A class for annotating labels on an image using provided detections. """ + @dataclass + class _TextProperties: + text: str + width: int + height: int + width_padded: int + height_padded: int + + _FONT = cv2.FONT_HERSHEY_SIMPLEX + def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.DEFAULT, @@ -1054,7 +1065,7 @@ def __init__( text_position: Position = Position.TOP_LEFT, color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, - use_smart_positioning: bool = False, + spread_out: bool = False, ): """ Args: @@ -1071,10 +1082,8 @@ def __init__( Options are `INDEX`, `CLASS`, `TRACK`. border_radius (int): The radius to apply round edges. If the selected value is higher than the lower dimension, width or height, is clipped. - use_smart_positioning (bool): Whether to use smart positioning to prevent - label overlapping or not. + spread_out (bool): Spread out the labels to avoid overlapping. """ - self.use_smart_positioning: bool = use_smart_positioning self.border_radius: int = border_radius self.color: Union[Color, ColorPalette] = color self.text_color: Union[Color, ColorPalette] = text_color @@ -1083,6 +1092,7 @@ def __init__( self.text_padding: int = text_padding self.text_anchor: Position = text_position self.color_lookup: ColorLookup = color_lookup + self.spread_out = spread_out @ensure_cv2_image_for_annotation def annotate( @@ -1136,23 +1146,23 @@ def annotate( assert isinstance(scene, np.ndarray) self._validate_labels(labels, detections) - # Get text properties for all detections - text_props = self._get_text_properties(detections, labels) + labels = self._get_labels_text(detections, labels) + text_properties = self._get_text_properties(labels) - # Calculate background coordinates for all labels - xyxy = self._calculate_label_backgrounds( - detections, text_props, self.text_anchor, self.text_padding + xyxy = self._calculate_label_positions( + detections, text_properties, self.text_anchor ) - # Adjust positions if smart positioning is enabled - if self.use_smart_positioning: - xyxy = spread_out(xyxy, step=2) + if self.spread_out: + xyxy = spread_out_boxes( + xyxy, + step=2, + max_iterations=len(xyxy) * 20) - # Draw all labels self._draw_labels( scene=scene, xyxy=xyxy, - text_props=text_props, + text_properties=text_properties, detections=detections, custom_color_lookup=custom_color_lookup, ) @@ -1160,7 +1170,6 @@ def annotate( return scene def _validate_labels(self, labels: Optional[List[str]], detections: Detections): - """Validates that the number of labels matches the number of detections.""" if labels is not None and len(labels) != len(detections): raise ValueError( f"The number of labels ({len(labels)}) does not match the " @@ -1168,57 +1177,51 @@ def _validate_labels(self, labels: Optional[List[str]], detections: Detections): f"should have exactly 1 label." ) - def _get_text_properties( - self, detections: Detections, custom_labels: Optional[List[str]] - ) -> List[dict]: + def _get_text_properties(self, labels: List[str]) -> List[_TextProperties]: """Gets text content and dimensions for all detections.""" - text_props = [] - font = cv2.FONT_HERSHEY_SIMPLEX - - for idx in range(len(detections)): - # Determine label text - text = self._get_label_text(detections, custom_labels, idx) + text_properties = [] - # Calculate text dimensions + for label in labels: (text_w, text_h) = cv2.getTextSize( - text=text, - fontFace=font, + text=label, + fontFace=self._FONT, fontScale=self.text_scale, thickness=self.text_thickness, )[0] - text_props.append( - { - "text": text, - "width": text_w, - "height": text_h, - "width_padded": text_w + 2 * self.text_padding, - "height_padded": text_h + 2 * self.text_padding, - } - ) + text_properties.append(self._TextProperties( + text=label, + width=text_w, + height=text_h, + width_padded=text_w + 2 * self.text_padding, + height_padded=text_h + 2 * self.text_padding, + )) - return text_props + return text_properties - def _get_label_text( - self, detections: Detections, custom_labels: Optional[List[str]], idx: int - ) -> str: - """Determines the label text for a given detection.""" + @staticmethod + def _get_labels_text( + detections: Detections, custom_labels: Optional[List[str]]) -> List[str]: + if custom_labels is not None: - return custom_labels[idx] - elif CLASS_NAME_DATA_FIELD in detections.data: - return detections.data[CLASS_NAME_DATA_FIELD][idx] - elif detections.class_id is not None: - return str(detections.class_id[idx]) - return str(idx) - - def _calculate_label_backgrounds( + return custom_labels + + labels = [] + for idx in range(len(detections)): + if CLASS_NAME_DATA_FIELD in detections.data: + labels.append(detections.data[CLASS_NAME_DATA_FIELD][idx]) + elif detections.class_id is not None: + labels.append(str(detections.class_id[idx])) + else: + labels.append(str(idx)) + return labels + + def _calculate_label_positions( self, detections: Detections, - text_props: List[dict], - text_anchor: str, - text_padding: int, + text_properties: List[_TextProperties], + text_anchor: Position, ) -> np.ndarray: - """Calculates background coordinates for all labels.""" anchors_coordinates = detections.get_anchors_coordinates( anchor=text_anchor ).astype(int) @@ -1228,8 +1231,8 @@ def _calculate_label_backgrounds( text_background_xyxy = resolve_text_background_xyxy( center_coordinates=tuple(center_coords), text_wh=( - text_props[idx]["width_padded"], - text_props[idx]["height_padded"], + text_properties[idx].width_padded, + text_properties[idx].height_padded, ), position=text_anchor, ) @@ -1239,22 +1242,24 @@ def _calculate_label_backgrounds( def _draw_labels( self, - scene: ImageType, + scene: np.ndarray, xyxy: np.ndarray, - text_props: List[dict], + text_properties: List[_TextProperties], detections: Detections, custom_color_lookup: Optional[np.ndarray], ) -> None: - """Draws all labels and their backgrounds on the scene.""" - if custom_color_lookup is not None: - color_lookup = custom_color_lookup - else: - color_lookup = self.color_lookup - font = cv2.FONT_HERSHEY_SIMPLEX + assert len(xyxy) == len(text_properties) == len(detections), ( + f"Number of text properties ({len(text_properties)}), xyxy ({len(xyxy)}) and detections ({len(detections)}) do not match." + ) - for idx, coordinates in enumerate(xyxy): - # Resolve colors - bg_color = resolve_color( + color_lookup = ( + custom_color_lookup + if custom_color_lookup is not None + else self.color_lookup + ) + + for idx, box_xyxy in enumerate(xyxy): + background_color = resolve_color( color=self.color, detections=detections, detection_idx=idx, @@ -1267,22 +1272,20 @@ def _draw_labels( color_lookup=color_lookup, ) - # Calculate text position - text_x = coordinates[0] + self.text_padding - text_y = coordinates[1] + self.text_padding + text_props[idx]["height"] - - # Draw background and text self.draw_rounded_rectangle( scene=scene, - xyxy=coordinates, - color=bg_color.as_bgr(), + xyxy=box_xyxy, + color=background_color.as_bgr(), border_radius=self.border_radius, ) + + text_x = box_xyxy[0] + self.text_padding + text_y = box_xyxy[1] + self.text_padding + text_properties[idx].height cv2.putText( img=scene, - text=text_props[idx]["text"], + text=text_properties[idx].text, org=(text_x, text_y), - fontFace=font, + fontFace=self._FONT, fontScale=self.text_scale, color=text_color.as_bgr(), thickness=self.text_thickness, @@ -1338,6 +1341,16 @@ class RichLabelAnnotator(BaseAnnotator): with support for Unicode characters by using a custom font. """ + @dataclass + class _TextProperties: + text: str + width: int + height: int + width_padded: int + height_padded: int + text_left: int + text_top: int + def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.DEFAULT, @@ -1348,7 +1361,7 @@ def __init__( text_position: Position = Position.TOP_LEFT, color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, - use_smart_positioning: bool = False, + spread_out: bool = False, ): """ Args: @@ -1365,8 +1378,7 @@ def __init__( Options are `INDEX`, `CLASS`, `TRACK`. border_radius (int): The radius to apply round edges. If the selected value is higher than the lower dimension, width or height, is clipped. - use_smart_positioning (bool): Whether to use smart positioning to prevent - label overlapping or not. + spread_out (bool): Spread out the labels to avoid overlapping. """ self.color = color self.text_color = text_color @@ -1374,15 +1386,8 @@ def __init__( self.text_anchor = text_position self.color_lookup = color_lookup self.border_radius = border_radius - self.use_smart_positioning: bool = use_smart_positioning - if font_path is not None: - try: - self.font = ImageFont.truetype(font_path, font_size) - except OSError: - print(f"Font path '{font_path}' not found. Using PIL's default font.") - self.font = self._load_default_font(font_size) - else: - self.font = self._load_default_font(font_size) + self.spread_out = spread_out + self.font = self._load_font(font_size, font_path) @ensure_pil_image_for_annotation def annotate( @@ -1432,152 +1437,185 @@ def annotate( """ assert isinstance(scene, Image.Image) + self._validate_labels(labels, detections) + draw = ImageDraw.Draw(scene) + labels = self._get_labels_text(detections, labels) + text_properties = self._get_text_properties(draw, labels) + + labels_text = self._get_labels_text(detections, labels) + xyxy = self._calculate_label_positions( + detections, text_properties, self.text_anchor + ) - # Input validation + if self.spread_out: + xyxy = spread_out_boxes( + xyxy, step=2, max_iterations=len(xyxy) * 20 + ) + + self._draw_labels( + draw=draw, + xyxy=xyxy, + labels=labels_text, + text_properties=text_properties, + detections=detections, + custom_color_lookup=custom_color_lookup + ) + + return scene + + def _validate_labels(self, labels: Optional[List[str]], detections: Detections): if labels is not None and len(labels) != len(detections): raise ValueError( - f"Label count ({len(labels)}) != detection count ({len(detections)})" + f"The number of labels ({len(labels)}) does not match the " + f"number of detections ({len(detections)}). Each detection " + f"should have exactly 1 label." ) - # Get anchor coordinates for all detections - detection_anchor_coordinates = detections.get_anchors_coordinates( + def _get_text_properties(self, draw, labels: List[str]) -> List[_TextProperties]: + """Gets text content and dimensions for all detections.""" + text_properties = [] + + for label in labels: + text_left, text_top, text_right, text_bottom = draw.textbbox( + (0, 0), label, font=self.font + ) + text_width = text_right - text_left + text_height = text_bottom - text_top + width_padded = text_width + 2 * self.text_padding + height_padded = text_height + 2 * self.text_padding + + text_properties.append(self._TextProperties( + text=label, + width=text_width, + height=text_height, + width_padded=width_padded, + height_padded=height_padded, + text_left=text_left, + text_top=text_top, + )) + + return text_properties + + def _calculate_label_positions(self, detections: Detections, text_properties: List[_TextProperties], text_anchor: Position) -> np.ndarray: + anchor_coordinates = detections.get_anchors_coordinates( anchor=self.text_anchor ).astype(int) - # Use the appropriate color lookup table - effective_color_lookup = ( + xyxy = [] + for idx, center_coords in enumerate(anchor_coordinates): + text_background_xyxy = resolve_text_background_xyxy( + center_coordinates=tuple(center_coords), + text_wh=( + text_properties[idx].width_padded, + text_properties[idx].height_padded, + ), + position=text_anchor, + ) + xyxy.append(text_background_xyxy) + + return np.array(xyxy) + + def _calculate_text_dimensions(self, draw, label_text: str) -> Tuple[Tuple[int, int], Tuple[int, int]]: + """ + Calculate text dimensions and offsets for the given label text. + Args: + label_text: The text to measure + Returns: + (Tuple[int, int]): ((left_offset, top_offset), (padded_width, padded_height)) + """ + text_left, text_top, text_right, text_bottom = draw.textbbox( + (0, 0), label_text, font=self.font + ) + text_width = text_right - text_left + text_height = text_bottom - text_top + padded_width = text_width + 2 * self.text_padding + padded_height = text_height + 2 * self.text_padding + return (text_left, text_top), (padded_width, padded_height) + + @staticmethod + def _get_labels_text( + detections: Detections, custom_labels: Optional[List[str]]) -> List[str]: + + if custom_labels is not None: + return custom_labels + + labels = [] + for idx in range(len(detections)): + if CLASS_NAME_DATA_FIELD in detections.data: + labels.append(detections.data[CLASS_NAME_DATA_FIELD][idx]) + elif detections.class_id is not None: + labels.append(str(detections.class_id[idx])) + else: + labels.append(str(idx)) + return labels + + def _draw_labels( + self, + draw, + xyxy: np.ndarray, + labels: List[str], + text_properties: List[_TextProperties], + detections: Detections, + custom_color_lookup: Optional[np.ndarray], + ) -> None: + color_lookup = ( custom_color_lookup if custom_color_lookup is not None else self.color_lookup ) - def _get_detection_label_text(detection_index: int) -> str: - """ - Determine the appropriate label text for a detection. - Args: - detection_index: Index of the current detection - Returns: - str: The label text to display - """ - if labels is not None: - return labels[detection_index] - if CLASS_NAME_DATA_FIELD in detections.data: - return detections.data[CLASS_NAME_DATA_FIELD][detection_index] - if detections.class_id is not None: - return str(detections.class_id[detection_index]) - return str(detection_index) - - def _calculate_text_dimensions(label_text: str) -> tuple: - """ - Calculate text dimensions and offsets for the given label text. - Args: - label_text: The text to measure - Returns: - tuple: ((left_offset, top_offset), (padded_width, padded_height)) - """ - text_left, text_top, text_right, text_bottom = draw.textbbox( - (0, 0), label_text, font=self.font - ) - text_width = text_right - text_left - text_height = text_bottom - text_top - padded_width = text_width + 2 * self.text_padding - padded_height = text_height + 2 * self.text_padding - return (text_left, text_top), (padded_width, padded_height) - - # Prepare all annotation data - annotation_collection = [] - for detection_index, center_coordinate in enumerate( - detection_anchor_coordinates - ): - # Get colors once per detection + for idx, box_xyxy in enumerate(xyxy): background_color = resolve_color( color=self.color, detections=detections, - detection_idx=detection_index, - color_lookup=effective_color_lookup, + detection_idx=idx, + color_lookup=color_lookup, ) - label_text_color = resolve_color( + text_color = resolve_color( color=self.text_color, detections=detections, - detection_idx=detection_index, - color_lookup=effective_color_lookup, - ) - - # Get text and calculate dimensions - label_text = _get_detection_label_text(detection_index) - text_offset_coordinates, padded_dimensions = _calculate_text_dimensions( - label_text - ) - - # Calculate background coordinates - background_coordinates = resolve_text_background_xyxy( - center_coordinates=tuple(center_coordinate), - text_wh=padded_dimensions, - position=self.text_anchor, - ) - - # Store all data for this annotation - annotation_collection.append( - { - "label_text": label_text, - "background_color": background_color, - "text_color": label_text_color, - "text_offset": text_offset_coordinates, - "background_coordinates": background_coordinates, - } - ) - - # Convert coordinates to numpy array for processing - background_coordinate_array = np.array( - [data["background_coordinates"] for data in annotation_collection] - ) - - # Apply smart positioning if enabled - if self.use_smart_positioning: - background_coordinate_array = spread_out( - background_coordinate_array, step=2 + detection_idx=idx, + color_lookup=color_lookup, ) - # Draw annotations - for annotation_index, coordinates in enumerate(background_coordinate_array): - annotation_data = annotation_collection[annotation_index] - - # Calculate final text position label_x_position = ( - coordinates[0] + self.text_padding - annotation_data["text_offset"][0] + box_xyxy[0] + self.text_padding - text_properties[idx].text_left ) label_y_position = ( - coordinates[1] + self.text_padding - annotation_data["text_offset"][1] + box_xyxy[1] + self.text_padding - text_properties[idx].text_top ) draw.rounded_rectangle( - tuple(coordinates), + tuple(box_xyxy), radius=self.border_radius, - fill=annotation_data["background_color"].as_rgb(), + fill=background_color.as_rgb(), outline=None, ) draw.text( xy=(label_x_position, label_y_position), - text=annotation_data["label_text"], + text=labels[idx], font=self.font, - fill=annotation_data["text_color"].as_rgb(), + fill=text_color.as_rgb(), ) - return scene + @staticmethod - def _load_default_font(size): - """ - PIL either loads a font that accepts a size (e.g. on my machine) - or raises an error saying `load_default` does not accept arguments - (e.g. in Colab). - """ + def _load_font(font_size: int, font_path: Optional[str]): + def load_default_font(size): + try: + return ImageFont.load_default(size) + except TypeError: + return ImageFont.load_default() + + if font_path is None: + return load_default_font(font_size) + try: - font = ImageFont.load_default(size) - except TypeError: - font = ImageFont.load_default() - return font - + return ImageFont.truetype(font_path, font_size) + except OSError: + print(f"Font path '{font_path}' not found. Using PIL's default font.") + return load_default_font(font_size) class IconAnnotator(BaseAnnotator): """ diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 308797eab..b874953df 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1042,83 +1042,53 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: return np.cross(vector_at_zero, anchors - vector_start) -# Intelligent padding functions -def get_intersection_center( +def get_box_intersection( xyxy_1: np.ndarray, xyxy_2: np.ndarray -) -> Optional[Tuple[float, float]]: +) -> Optional[np.ndarray]: overlap_xmin = max(xyxy_1[0], xyxy_2[0]) overlap_ymin = max(xyxy_1[1], xyxy_2[1]) overlap_xmax = min(xyxy_1[2], xyxy_2[2]) overlap_ymax = min(xyxy_1[3], xyxy_2[3]) if overlap_xmin < overlap_xmax and overlap_ymin < overlap_ymax: - x_center = (overlap_xmin + overlap_xmax) / 2 - y_center = (overlap_ymin + overlap_ymax) / 2 - return (x_center, y_center) + return np.array([overlap_xmin, overlap_ymin, overlap_xmax, overlap_ymax]) else: return None +def get_unit_vector(xy_1: np.ndarray, xy_2: np.ndarray) -> np.ndarray: + direction = xy_2 - xy_1 + magnitude = np.linalg.norm(direction) + return direction / magnitude if magnitude > 0 else np.zeros(2) -def get_box_center(xyxy: np.ndarray) -> Tuple[float, float]: - x_center = (xyxy[0] + xyxy[2]) / 2 - y_center = (xyxy[1] + xyxy[3]) / 2 - return (x_center, y_center) +def spread_out_boxes(xyxy: np.ndarray, step: int, max_iterations: int = 100) -> np.ndarray: + if len(xyxy) == 0: + return xyxy - -def vector_with_length( - xy_1: Tuple[float, float], xy_2: Tuple[float, float], n: float -) -> Tuple[float, float]: - x1, y1 = xy_1 - x2, y2 = xy_2 - - dx = x2 - x1 - dy = y2 - y1 - - if dx == 0 and dy == 0: - return 0, 0 - - magnitude = math.sqrt(dx**2 + dy**2) - - unit_dx = dx / magnitude - unit_dy = dy / magnitude - - v1 = unit_dx * n - v2 = unit_dy * n - - return (v1, v2) - - -def pad(xyxy: np.ndarray, px: int, py: Optional[int] = None): - if py is None: - py = px - - result = xyxy.copy() - result[:, [0, 1]] -= [px, py] - result[:, [2, 3]] += [px, py] - - return result - - -def spread_out(xyxy: np.ndarray, step) -> np.ndarray: - xyxy_padded = pad(xyxy, px=step) - while True: + xyxy_padded = pad_boxes(xyxy, px=step) + for _ in range(max_iterations): iou = box_iou_batch(xyxy_padded, xyxy_padded) np.fill_diagonal(iou, 0) if np.all(iou == 0): - return pad(xyxy_padded, px=-step) + break i, j = np.unravel_index(np.argmax(iou), iou.shape) xyxy_i, xyxy_j = xyxy_padded[i], xyxy_padded[j] - intersection_center = get_intersection_center(xyxy_i, xyxy_j) - xyxy_i_center = get_box_center(xyxy_i) - xyxy_j_center = get_box_center(xyxy_j) + box_intersection = get_box_intersection(xyxy_i, xyxy_j) + assert box_intersection is not None, \ + "Since we checked IoU already, boxes should always intersect" + + intersection_center = (box_intersection[:2] + box_intersection[2:]) / 2 + xyxy_i_center = (xyxy_i[:2] + xyxy_i[2:]) / 2 + xyxy_j_center = (xyxy_j[:2] + xyxy_j[2:]) / 2 + + unit_vector_i = get_unit_vector(intersection_center, xyxy_i_center) + unit_vector_j = get_unit_vector(intersection_center, xyxy_j_center) - vector_i = vector_with_length(intersection_center, xyxy_i_center, step) - vector_j = vector_with_length(intersection_center, xyxy_j_center, step) + xyxy_padded[i, [0, 2]] += int(unit_vector_i[0] * step) + xyxy_padded[i, [1, 3]] += int(unit_vector_i[1] * step) + xyxy_padded[j, [0, 2]] += int(unit_vector_j[0] * step) + xyxy_padded[j, [1, 3]] += int(unit_vector_j[1] * step) - xyxy_padded[i, [0, 2]] += int(vector_i[0]) - xyxy_padded[i, [1, 3]] += int(vector_i[1]) - xyxy_padded[j, [0, 2]] += int(vector_j[0]) - xyxy_padded[j, [1, 3]] += int(vector_j[1]) + return pad_boxes(xyxy_padded, px=-step) diff --git a/supervision/keypoint/annotators.py b/supervision/keypoint/annotators.py index e82acbc65..ba1c5448d 100644 --- a/supervision/keypoint/annotators.py +++ b/supervision/keypoint/annotators.py @@ -5,9 +5,9 @@ import cv2 import numpy as np -from supervision import Rect, pad_boxes +from supervision.geometry.core import Rect from supervision.annotators.base import ImageType -from supervision.detection.utils import pad, spread_out +from supervision.detection.utils import pad_boxes, spread_out_boxes from supervision.draw.color import Color from supervision.draw.utils import draw_rounded_rectangle from supervision.keypoint.core import KeyPoints @@ -202,7 +202,7 @@ def __init__( text_thickness: int = 1, text_padding: int = 10, border_radius: int = 0, - use_smart_positioning: bool = False, + spread_out: bool = False, ): """ Args: @@ -217,16 +217,15 @@ def __init__( text_padding (int): The padding around the text. border_radius (int): The radius of the rounded corners of the boxes. Set to a high value to produce circles. - use_smart_positioning (bool): Whether to use smart positioning to prevent - label overlapping or not. + spread_out (bool): Spread out the labels to avoid overlap. """ - self.use_smart_positioning = use_smart_positioning self.border_radius: int = border_radius self.color: Union[Color, List[Color]] = color self.text_color: Union[Color, List[Color]] = text_color self.text_scale: float = text_scale self.text_thickness: int = text_thickness self.text_padding: int = text_padding + self.spread_out = spread_out def annotate( self, @@ -349,25 +348,25 @@ def annotate( text_colors = text_colors[mask] labels = labels[mask] - xyxy = np.array( - [ - self.get_text_bounding_box( - text=label, - font=font, - text_scale=self.text_scale, - text_thickness=self.text_thickness, - center_coordinates=tuple(anchor), - ) - for anchor, label in zip(anchors, labels) - ] - ) - - if self.use_smart_positioning: - xyxy_padded = pad(xyxy=xyxy, px=self.text_padding) - xyxy_padded = spread_out(xyxy_padded, step=2) - xyxy = pad(xyxy=xyxy_padded, px=-self.text_padding) - else: - xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding) + xyxy = np.array([ + self.get_text_bounding_box( + text=label, + font=font, + text_scale=self.text_scale, + text_thickness=self.text_thickness, + center_coordinates=tuple(anchor), + ) + for anchor, label in zip(anchors, labels) + ]) + xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding) + + if self.spread_out: + xyxy_padded = spread_out_boxes( + xyxy_padded, + step=2, + max_iterations=len(xyxy_padded) * 20 + ) + xyxy = pad_boxes(xyxy=xyxy_padded, px=-self.text_padding) for text, color, text_color, box, box_padded in zip( labels, colors, text_colors, xyxy, xyxy_padded From 920a2c0cda8251872339d49dc59180e27c50b20d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 7 Nov 2024 15:27:30 +0000 Subject: [PATCH 05/16] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto?= =?UTF-8?q?=20format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/annotators/core.py | 86 ++++++++++++++++-------------- supervision/detection/utils.py | 12 +++-- supervision/keypoint/annotators.py | 28 +++++----- 3 files changed, 68 insertions(+), 58 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 2864bfdd9..144e7c7f8 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,7 +1,7 @@ from dataclasses import dataclass from functools import lru_cache from math import sqrt -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import cv2 import numpy as np @@ -1154,10 +1154,7 @@ def annotate( ) if self.spread_out: - xyxy = spread_out_boxes( - xyxy, - step=2, - max_iterations=len(xyxy) * 20) + xyxy = spread_out_boxes(xyxy, step=2, max_iterations=len(xyxy) * 20) self._draw_labels( scene=scene, @@ -1189,23 +1186,25 @@ def _get_text_properties(self, labels: List[str]) -> List[_TextProperties]: thickness=self.text_thickness, )[0] - text_properties.append(self._TextProperties( - text=label, - width=text_w, - height=text_h, - width_padded=text_w + 2 * self.text_padding, - height_padded=text_h + 2 * self.text_padding, - )) + text_properties.append( + self._TextProperties( + text=label, + width=text_w, + height=text_h, + width_padded=text_w + 2 * self.text_padding, + height_padded=text_h + 2 * self.text_padding, + ) + ) return text_properties @staticmethod def _get_labels_text( - detections: Detections, custom_labels: Optional[List[str]]) -> List[str]: - + detections: Detections, custom_labels: Optional[List[str]] + ) -> List[str]: if custom_labels is not None: return custom_labels - + labels = [] for idx in range(len(detections)): if CLASS_NAME_DATA_FIELD in detections.data: @@ -1248,9 +1247,9 @@ def _draw_labels( detections: Detections, custom_color_lookup: Optional[np.ndarray], ) -> None: - assert len(xyxy) == len(text_properties) == len(detections), ( - f"Number of text properties ({len(text_properties)}), xyxy ({len(xyxy)}) and detections ({len(detections)}) do not match." - ) + assert ( + len(xyxy) == len(text_properties) == len(detections) + ), f"Number of text properties ({len(text_properties)}), xyxy ({len(xyxy)}) and detections ({len(detections)}) do not match." color_lookup = ( custom_color_lookup @@ -1449,9 +1448,7 @@ def annotate( ) if self.spread_out: - xyxy = spread_out_boxes( - xyxy, step=2, max_iterations=len(xyxy) * 20 - ) + xyxy = spread_out_boxes(xyxy, step=2, max_iterations=len(xyxy) * 20) self._draw_labels( draw=draw, @@ -1459,7 +1456,7 @@ def annotate( labels=labels_text, text_properties=text_properties, detections=detections, - custom_color_lookup=custom_color_lookup + custom_color_lookup=custom_color_lookup, ) return scene @@ -1485,19 +1482,26 @@ def _get_text_properties(self, draw, labels: List[str]) -> List[_TextProperties] width_padded = text_width + 2 * self.text_padding height_padded = text_height + 2 * self.text_padding - text_properties.append(self._TextProperties( - text=label, - width=text_width, - height=text_height, - width_padded=width_padded, - height_padded=height_padded, - text_left=text_left, - text_top=text_top, - )) + text_properties.append( + self._TextProperties( + text=label, + width=text_width, + height=text_height, + width_padded=width_padded, + height_padded=height_padded, + text_left=text_left, + text_top=text_top, + ) + ) return text_properties - - def _calculate_label_positions(self, detections: Detections, text_properties: List[_TextProperties], text_anchor: Position) -> np.ndarray: + + def _calculate_label_positions( + self, + detections: Detections, + text_properties: List[_TextProperties], + text_anchor: Position, + ) -> np.ndarray: anchor_coordinates = detections.get_anchors_coordinates( anchor=self.text_anchor ).astype(int) @@ -1516,7 +1520,9 @@ def _calculate_label_positions(self, detections: Detections, text_properties: Li return np.array(xyxy) - def _calculate_text_dimensions(self, draw, label_text: str) -> Tuple[Tuple[int, int], Tuple[int, int]]: + def _calculate_text_dimensions( + self, draw, label_text: str + ) -> Tuple[Tuple[int, int], Tuple[int, int]]: """ Calculate text dimensions and offsets for the given label text. Args: @@ -1535,11 +1541,11 @@ def _calculate_text_dimensions(self, draw, label_text: str) -> Tuple[Tuple[int, @staticmethod def _get_labels_text( - detections: Detections, custom_labels: Optional[List[str]]) -> List[str]: - + detections: Detections, custom_labels: Optional[List[str]] + ) -> List[str]: if custom_labels is not None: return custom_labels - + labels = [] for idx in range(len(detections)): if CLASS_NAME_DATA_FIELD in detections.data: @@ -1599,7 +1605,6 @@ def _draw_labels( fill=text_color.as_rgb(), ) - @staticmethod def _load_font(font_size: int, font_path: Optional[str]): def load_default_font(size): @@ -1607,16 +1612,17 @@ def load_default_font(size): return ImageFont.load_default(size) except TypeError: return ImageFont.load_default() - + if font_path is None: return load_default_font(font_size) - + try: return ImageFont.truetype(font_path, font_size) except OSError: print(f"Font path '{font_path}' not found. Using PIL's default font.") return load_default_font(font_size) + class IconAnnotator(BaseAnnotator): """ A class for drawing an icon on an image, using provided detections. diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index b874953df..ede1ed355 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1,4 +1,3 @@ -import math from itertools import chain from typing import Dict, List, Optional, Tuple, Union @@ -1055,12 +1054,16 @@ def get_box_intersection( else: return None + def get_unit_vector(xy_1: np.ndarray, xy_2: np.ndarray) -> np.ndarray: direction = xy_2 - xy_1 magnitude = np.linalg.norm(direction) return direction / magnitude if magnitude > 0 else np.zeros(2) -def spread_out_boxes(xyxy: np.ndarray, step: int, max_iterations: int = 100) -> np.ndarray: + +def spread_out_boxes( + xyxy: np.ndarray, step: int, max_iterations: int = 100 +) -> np.ndarray: if len(xyxy) == 0: return xyxy @@ -1076,8 +1079,9 @@ def spread_out_boxes(xyxy: np.ndarray, step: int, max_iterations: int = 100) -> xyxy_i, xyxy_j = xyxy_padded[i], xyxy_padded[j] box_intersection = get_box_intersection(xyxy_i, xyxy_j) - assert box_intersection is not None, \ - "Since we checked IoU already, boxes should always intersect" + assert ( + box_intersection is not None + ), "Since we checked IoU already, boxes should always intersect" intersection_center = (box_intersection[:2] + box_intersection[2:]) / 2 xyxy_i_center = (xyxy_i[:2] + xyxy_i[2:]) / 2 diff --git a/supervision/keypoint/annotators.py b/supervision/keypoint/annotators.py index ba1c5448d..4d3196b39 100644 --- a/supervision/keypoint/annotators.py +++ b/supervision/keypoint/annotators.py @@ -5,11 +5,11 @@ import cv2 import numpy as np -from supervision.geometry.core import Rect from supervision.annotators.base import ImageType from supervision.detection.utils import pad_boxes, spread_out_boxes from supervision.draw.color import Color from supervision.draw.utils import draw_rounded_rectangle +from supervision.geometry.core import Rect from supervision.keypoint.core import KeyPoints from supervision.keypoint.skeletons import SKELETONS_BY_VERTEX_COUNT from supervision.utils.conversion import ensure_cv2_image_for_annotation @@ -348,23 +348,23 @@ def annotate( text_colors = text_colors[mask] labels = labels[mask] - xyxy = np.array([ - self.get_text_bounding_box( - text=label, - font=font, - text_scale=self.text_scale, - text_thickness=self.text_thickness, - center_coordinates=tuple(anchor), - ) - for anchor, label in zip(anchors, labels) - ]) + xyxy = np.array( + [ + self.get_text_bounding_box( + text=label, + font=font, + text_scale=self.text_scale, + text_thickness=self.text_thickness, + center_coordinates=tuple(anchor), + ) + for anchor, label in zip(anchors, labels) + ] + ) xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding) if self.spread_out: xyxy_padded = spread_out_boxes( - xyxy_padded, - step=2, - max_iterations=len(xyxy_padded) * 20 + xyxy_padded, step=2, max_iterations=len(xyxy_padded) * 20 ) xyxy = pad_boxes(xyxy=xyxy_padded, px=-self.text_padding) From b66702c5ffae69ecf51bb1c6b9da44dce39b5af8 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Thu, 7 Nov 2024 17:32:09 +0200 Subject: [PATCH 06/16] Remove dead code, surplus function args --- supervision/annotators/core.py | 32 ++++++-------------------------- 1 file changed, 6 insertions(+), 26 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 144e7c7f8..bfc20503b 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1247,9 +1247,11 @@ def _draw_labels( detections: Detections, custom_color_lookup: Optional[np.ndarray], ) -> None: - assert ( - len(xyxy) == len(text_properties) == len(detections) - ), f"Number of text properties ({len(text_properties)}), xyxy ({len(xyxy)}) and detections ({len(detections)}) do not match." + assert len(xyxy) == len(text_properties) == len(detections), ( + f"Number of text properties ({len(text_properties)}), " + f"xyxy ({len(xyxy)}) and detections ({len(detections)}) " + "do not match." + ) color_lookup = ( custom_color_lookup @@ -1442,7 +1444,6 @@ def annotate( labels = self._get_labels_text(detections, labels) text_properties = self._get_text_properties(draw, labels) - labels_text = self._get_labels_text(detections, labels) xyxy = self._calculate_label_positions( detections, text_properties, self.text_anchor ) @@ -1453,7 +1454,6 @@ def annotate( self._draw_labels( draw=draw, xyxy=xyxy, - labels=labels_text, text_properties=text_properties, detections=detections, custom_color_lookup=custom_color_lookup, @@ -1520,25 +1520,6 @@ def _calculate_label_positions( return np.array(xyxy) - def _calculate_text_dimensions( - self, draw, label_text: str - ) -> Tuple[Tuple[int, int], Tuple[int, int]]: - """ - Calculate text dimensions and offsets for the given label text. - Args: - label_text: The text to measure - Returns: - (Tuple[int, int]): ((left_offset, top_offset), (padded_width, padded_height)) - """ - text_left, text_top, text_right, text_bottom = draw.textbbox( - (0, 0), label_text, font=self.font - ) - text_width = text_right - text_left - text_height = text_bottom - text_top - padded_width = text_width + 2 * self.text_padding - padded_height = text_height + 2 * self.text_padding - return (text_left, text_top), (padded_width, padded_height) - @staticmethod def _get_labels_text( detections: Detections, custom_labels: Optional[List[str]] @@ -1560,7 +1541,6 @@ def _draw_labels( self, draw, xyxy: np.ndarray, - labels: List[str], text_properties: List[_TextProperties], detections: Detections, custom_color_lookup: Optional[np.ndarray], @@ -1600,7 +1580,7 @@ def _draw_labels( ) draw.text( xy=(label_x_position, label_y_position), - text=labels[idx], + text=text_properties[idx].text, font=self.font, fill=text_color.as_rgb(), ) From 5fedfe522087d55ee5bf081c59fc912e785b7a90 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Fri, 8 Nov 2024 13:07:34 +0200 Subject: [PATCH 07/16] Rename "spread_out" arg to "smart_positions" --- supervision/annotators/core.py | 16 ++++++++-------- supervision/keypoint/annotators.py | 8 ++++---- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index bfc20503b..44b816457 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1065,7 +1065,7 @@ def __init__( text_position: Position = Position.TOP_LEFT, color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, - spread_out: bool = False, + smart_positions: bool = False, ): """ Args: @@ -1082,7 +1082,7 @@ def __init__( Options are `INDEX`, `CLASS`, `TRACK`. border_radius (int): The radius to apply round edges. If the selected value is higher than the lower dimension, width or height, is clipped. - spread_out (bool): Spread out the labels to avoid overlapping. + smart_positions (bool): Spread out the labels to avoid overlapping. """ self.border_radius: int = border_radius self.color: Union[Color, ColorPalette] = color @@ -1092,7 +1092,7 @@ def __init__( self.text_padding: int = text_padding self.text_anchor: Position = text_position self.color_lookup: ColorLookup = color_lookup - self.spread_out = spread_out + self.smart_positions = smart_positions @ensure_cv2_image_for_annotation def annotate( @@ -1153,7 +1153,7 @@ def annotate( detections, text_properties, self.text_anchor ) - if self.spread_out: + if self.smart_positions: xyxy = spread_out_boxes(xyxy, step=2, max_iterations=len(xyxy) * 20) self._draw_labels( @@ -1362,7 +1362,7 @@ def __init__( text_position: Position = Position.TOP_LEFT, color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, - spread_out: bool = False, + smart_positions: bool = False, ): """ Args: @@ -1379,7 +1379,7 @@ def __init__( Options are `INDEX`, `CLASS`, `TRACK`. border_radius (int): The radius to apply round edges. If the selected value is higher than the lower dimension, width or height, is clipped. - spread_out (bool): Spread out the labels to avoid overlapping. + smart_positions (bool): Spread out the labels to avoid overlapping. """ self.color = color self.text_color = text_color @@ -1387,7 +1387,7 @@ def __init__( self.text_anchor = text_position self.color_lookup = color_lookup self.border_radius = border_radius - self.spread_out = spread_out + self.smart_positions = smart_positions self.font = self._load_font(font_size, font_path) @ensure_pil_image_for_annotation @@ -1448,7 +1448,7 @@ def annotate( detections, text_properties, self.text_anchor ) - if self.spread_out: + if self.smart_positions: xyxy = spread_out_boxes(xyxy, step=2, max_iterations=len(xyxy) * 20) self._draw_labels( diff --git a/supervision/keypoint/annotators.py b/supervision/keypoint/annotators.py index 4d3196b39..d968786c8 100644 --- a/supervision/keypoint/annotators.py +++ b/supervision/keypoint/annotators.py @@ -202,7 +202,7 @@ def __init__( text_thickness: int = 1, text_padding: int = 10, border_radius: int = 0, - spread_out: bool = False, + smart_positions: bool = False, ): """ Args: @@ -217,7 +217,7 @@ def __init__( text_padding (int): The padding around the text. border_radius (int): The radius of the rounded corners of the boxes. Set to a high value to produce circles. - spread_out (bool): Spread out the labels to avoid overlap. + smart_positions (bool): Spread out the labels to avoid overlap. """ self.border_radius: int = border_radius self.color: Union[Color, List[Color]] = color @@ -225,7 +225,7 @@ def __init__( self.text_scale: float = text_scale self.text_thickness: int = text_thickness self.text_padding: int = text_padding - self.spread_out = spread_out + self.smart_positions = smart_positions def annotate( self, @@ -362,7 +362,7 @@ def annotate( ) xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding) - if self.spread_out: + if self.smart_positions: xyxy_padded = spread_out_boxes( xyxy_padded, step=2, max_iterations=len(xyxy_padded) * 20 ) From 01c2912a64426d948a36f372e83f4662796e5f35 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Fri, 8 Nov 2024 13:10:18 +0200 Subject: [PATCH 08/16] Move FONT to global scope --- supervision/annotators/core.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 44b816457..6071b0839 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -33,6 +33,8 @@ ) from supervision.utils.internal import deprecated +CV2_FONT = cv2.FONT_HERSHEY_SIMPLEX + class BoxAnnotator(BaseAnnotator): """ @@ -1053,8 +1055,6 @@ class _TextProperties: width_padded: int height_padded: int - _FONT = cv2.FONT_HERSHEY_SIMPLEX - def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.DEFAULT, @@ -1181,7 +1181,7 @@ def _get_text_properties(self, labels: List[str]) -> List[_TextProperties]: for label in labels: (text_w, text_h) = cv2.getTextSize( text=label, - fontFace=self._FONT, + fontFace=CV2_FONT, fontScale=self.text_scale, thickness=self.text_thickness, )[0] @@ -1286,7 +1286,7 @@ def _draw_labels( img=scene, text=text_properties[idx].text, org=(text_x, text_y), - fontFace=self._FONT, + fontFace=CV2_FONT, fontScale=self.text_scale, color=text_color.as_bgr(), thickness=self.text_thickness, From 5e76f86ff938d018da49067981b182f6a3cfd5c4 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Fri, 8 Nov 2024 15:47:37 +0200 Subject: [PATCH 09/16] 3 changes to spread_put_boxes algo: * Vectorized * Using forces rather than discrete steps * Move along secondary axis more --- supervision/annotators/core.py | 4 +-- supervision/detection/utils.py | 53 +++++++++++++++++++----------- supervision/keypoint/annotators.py | 4 +-- 3 files changed, 37 insertions(+), 24 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 6071b0839..3200dc9c2 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1154,7 +1154,7 @@ def annotate( ) if self.smart_positions: - xyxy = spread_out_boxes(xyxy, step=2, max_iterations=len(xyxy) * 20) + xyxy = spread_out_boxes(xyxy) self._draw_labels( scene=scene, @@ -1449,7 +1449,7 @@ def annotate( ) if self.smart_positions: - xyxy = spread_out_boxes(xyxy, step=2, max_iterations=len(xyxy) * 20) + xyxy = spread_out_boxes(xyxy) self._draw_labels( draw=draw, diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index ede1ed355..246d3c607 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1062,37 +1062,52 @@ def get_unit_vector(xy_1: np.ndarray, xy_2: np.ndarray) -> np.ndarray: def spread_out_boxes( - xyxy: np.ndarray, step: int, max_iterations: int = 100 + xyxy: np.ndarray, + max_iterations: int = 100, + force_multiplier: float = 0.03, ) -> np.ndarray: + """ + Spread out boxes that overlap with each other. + + Args: + xyxy: Numpy array of shape (N, 4) where N is the number of boxes. + max_iterations: Maximum number of iterations to run the algorithm for. + force_multiplier: Multiplier to scale the force vectors by. Similar to + learning rate in gradient descent. + """ if len(xyxy) == 0: return xyxy - xyxy_padded = pad_boxes(xyxy, px=step) + xyxy_padded = pad_boxes(xyxy, px=1) for _ in range(max_iterations): + # NxN iou = box_iou_batch(xyxy_padded, xyxy_padded) np.fill_diagonal(iou, 0) - if np.all(iou == 0): break - i, j = np.unravel_index(np.argmax(iou), iou.shape) + overlap_mask = iou > 0 + + # Nx2 + centers = (xyxy_padded[:, :2] + xyxy_padded[:, 2:]) / 2 + + # NxNx2 + delta_centers = centers[:, np.newaxis, :] - centers[np.newaxis, :, :] + delta_centers *= overlap_mask[:, :, np.newaxis] - xyxy_i, xyxy_j = xyxy_padded[i], xyxy_padded[j] - box_intersection = get_box_intersection(xyxy_i, xyxy_j) - assert ( - box_intersection is not None - ), "Since we checked IoU already, boxes should always intersect" + # Nx2 + force_vectors = np.sum(delta_centers, axis=1) + force_vectors *= force_multiplier + force_vectors[(force_vectors > 0) & (force_vectors < 1)] = 1 + force_vectors[(force_vectors < 0) & (force_vectors > -1)] = -1 - intersection_center = (box_intersection[:2] + box_intersection[2:]) / 2 - xyxy_i_center = (xyxy_i[:2] + xyxy_i[2:]) / 2 - xyxy_j_center = (xyxy_j[:2] + xyxy_j[2:]) / 2 + # Reduce motion along primary axis + primary_axis = np.argmax(np.abs(force_vectors), axis=1) + force_vectors[np.arange(len(force_vectors)), primary_axis] /= 2 - unit_vector_i = get_unit_vector(intersection_center, xyxy_i_center) - unit_vector_j = get_unit_vector(intersection_center, xyxy_j_center) + force_vectors = force_vectors.astype(int) - xyxy_padded[i, [0, 2]] += int(unit_vector_i[0] * step) - xyxy_padded[i, [1, 3]] += int(unit_vector_i[1] * step) - xyxy_padded[j, [0, 2]] += int(unit_vector_j[0] * step) - xyxy_padded[j, [1, 3]] += int(unit_vector_j[1] * step) + xyxy_padded[:, [0, 1]] += force_vectors + xyxy_padded[:, [2, 3]] += force_vectors - return pad_boxes(xyxy_padded, px=-step) + return pad_boxes(xyxy_padded, px=-1) diff --git a/supervision/keypoint/annotators.py b/supervision/keypoint/annotators.py index d968786c8..fc450a837 100644 --- a/supervision/keypoint/annotators.py +++ b/supervision/keypoint/annotators.py @@ -363,9 +363,7 @@ def annotate( xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding) if self.smart_positions: - xyxy_padded = spread_out_boxes( - xyxy_padded, step=2, max_iterations=len(xyxy_padded) * 20 - ) + xyxy_padded = spread_out_boxes(xyxy_padded) xyxy = pad_boxes(xyxy=xyxy_padded, px=-self.text_padding) for text, color, text_color, box, box_padded in zip( From 38dbf5d05dd4a84cadc826c9a91e8575a4a6e57f Mon Sep 17 00:00:00 2001 From: LinasKo Date: Fri, 8 Nov 2024 16:42:04 +0200 Subject: [PATCH 10/16] Remove TextProperties class, merge auxiliary functions, use numpy arrays for passing data --- supervision/annotators/core.py | 198 ++++++++++++--------------------- 1 file changed, 73 insertions(+), 125 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 3200dc9c2..f8ab90827 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,4 +1,3 @@ -from dataclasses import dataclass from functools import lru_cache from math import sqrt from typing import List, Optional, Tuple, Union @@ -1047,14 +1046,6 @@ class LabelAnnotator(BaseAnnotator): A class for annotating labels on an image using provided detections. """ - @dataclass - class _TextProperties: - text: str - width: int - height: int - width_padded: int - height_padded: int - def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.DEFAULT, @@ -1147,19 +1138,17 @@ def annotate( self._validate_labels(labels, detections) labels = self._get_labels_text(detections, labels) - text_properties = self._get_text_properties(labels) - - xyxy = self._calculate_label_positions( - detections, text_properties, self.text_anchor - ) + label_properties = self._get_label_properties(detections, labels) + xyxy = label_properties[:, :4] if self.smart_positions: xyxy = spread_out_boxes(xyxy) + label_properties[:, :4] = xyxy self._draw_labels( scene=scene, - xyxy=xyxy, - text_properties=text_properties, + labels=labels, + label_properties=label_properties, detections=detections, custom_color_lookup=custom_color_lookup, ) @@ -1174,11 +1163,17 @@ def _validate_labels(self, labels: Optional[List[str]], detections: Detections): f"should have exactly 1 label." ) - def _get_text_properties(self, labels: List[str]) -> List[_TextProperties]: - """Gets text content and dimensions for all detections.""" - text_properties = [] + def _get_label_properties( + self, + detections: Detections, + labels: List[str], + ) -> np.ndarray: + label_properties = [] + anchors_coordinates = detections.get_anchors_coordinates( + anchor=self.text_anchor + ).astype(int) - for label in labels: + for label, center_coords in zip(labels, anchors_coordinates): (text_w, text_h) = cv2.getTextSize( text=label, fontFace=CV2_FONT, @@ -1186,17 +1181,23 @@ def _get_text_properties(self, labels: List[str]) -> List[_TextProperties]: thickness=self.text_thickness, )[0] - text_properties.append( - self._TextProperties( - text=label, - width=text_w, - height=text_h, - width_padded=text_w + 2 * self.text_padding, - height_padded=text_h + 2 * self.text_padding, - ) + width_padded = text_w + 2 * self.text_padding + height_padded = text_h + 2 * self.text_padding + + text_background_xyxy = resolve_text_background_xyxy( + center_coordinates=tuple(center_coords), + text_wh=(width_padded, height_padded), + position=self.text_anchor, ) - return text_properties + label_properties.append( + [ + *text_background_xyxy, + text_h, + ] + ) + + return np.array(label_properties).reshape(-1, 5) @staticmethod def _get_labels_text( @@ -1215,41 +1216,17 @@ def _get_labels_text( labels.append(str(idx)) return labels - def _calculate_label_positions( - self, - detections: Detections, - text_properties: List[_TextProperties], - text_anchor: Position, - ) -> np.ndarray: - anchors_coordinates = detections.get_anchors_coordinates( - anchor=text_anchor - ).astype(int) - - xyxy = [] - for idx, center_coords in enumerate(anchors_coordinates): - text_background_xyxy = resolve_text_background_xyxy( - center_coordinates=tuple(center_coords), - text_wh=( - text_properties[idx].width_padded, - text_properties[idx].height_padded, - ), - position=text_anchor, - ) - xyxy.append(text_background_xyxy) - - return np.array(xyxy) - def _draw_labels( self, scene: np.ndarray, - xyxy: np.ndarray, - text_properties: List[_TextProperties], + labels: List[str], + label_properties: np.ndarray, detections: Detections, custom_color_lookup: Optional[np.ndarray], ) -> None: - assert len(xyxy) == len(text_properties) == len(detections), ( - f"Number of text properties ({len(text_properties)}), " - f"xyxy ({len(xyxy)}) and detections ({len(detections)}) " + assert len(labels) == len(label_properties) == len(detections), ( + f"Number of label properties ({len(label_properties)}), " + f"labels ({len(labels)}) and detections ({len(detections)}) " "do not match." ) @@ -1259,7 +1236,7 @@ def _draw_labels( else self.color_lookup ) - for idx, box_xyxy in enumerate(xyxy): + for idx, label_property in enumerate(label_properties): background_color = resolve_color( color=self.color, detections=detections, @@ -1273,6 +1250,8 @@ def _draw_labels( color_lookup=color_lookup, ) + box_xyxy = label_property[:4] + text_height_padded = label_property[4] self.draw_rounded_rectangle( scene=scene, xyxy=box_xyxy, @@ -1281,10 +1260,10 @@ def _draw_labels( ) text_x = box_xyxy[0] + self.text_padding - text_y = box_xyxy[1] + self.text_padding + text_properties[idx].height + text_y = box_xyxy[1] + self.text_padding + text_height_padded cv2.putText( img=scene, - text=text_properties[idx].text, + text=labels[idx], org=(text_x, text_y), fontFace=CV2_FONT, fontScale=self.text_scale, @@ -1342,16 +1321,6 @@ class RichLabelAnnotator(BaseAnnotator): with support for Unicode characters by using a custom font. """ - @dataclass - class _TextProperties: - text: str - width: int - height: int - width_padded: int - height_padded: int - text_left: int - text_top: int - def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.DEFAULT, @@ -1442,19 +1411,17 @@ def annotate( draw = ImageDraw.Draw(scene) labels = self._get_labels_text(detections, labels) - text_properties = self._get_text_properties(draw, labels) - - xyxy = self._calculate_label_positions( - detections, text_properties, self.text_anchor - ) + label_properties = self._get_label_properties(draw, detections, labels) + xyxy = label_properties[:, :4] if self.smart_positions: xyxy = spread_out_boxes(xyxy) + label_properties[:, :4] = xyxy self._draw_labels( draw=draw, - xyxy=xyxy, - text_properties=text_properties, + labels=labels, + label_properties=label_properties, detections=detections, custom_color_lookup=custom_color_lookup, ) @@ -1469,11 +1436,16 @@ def _validate_labels(self, labels: Optional[List[str]], detections: Detections): f"should have exactly 1 label." ) - def _get_text_properties(self, draw, labels: List[str]) -> List[_TextProperties]: - """Gets text content and dimensions for all detections.""" - text_properties = [] + def _get_label_properties( + self, draw, detections: Detections, labels: List[str] + ) -> np.ndarray: + label_properties = [] - for label in labels: + anchor_coordinates = detections.get_anchors_coordinates( + anchor=self.text_anchor + ).astype(int) + + for label, center_coords in zip(labels, anchor_coordinates): text_left, text_top, text_right, text_bottom = draw.textbbox( (0, 0), label, font=self.font ) @@ -1482,43 +1454,15 @@ def _get_text_properties(self, draw, labels: List[str]) -> List[_TextProperties] width_padded = text_width + 2 * self.text_padding height_padded = text_height + 2 * self.text_padding - text_properties.append( - self._TextProperties( - text=label, - width=text_width, - height=text_height, - width_padded=width_padded, - height_padded=height_padded, - text_left=text_left, - text_top=text_top, - ) - ) - - return text_properties - - def _calculate_label_positions( - self, - detections: Detections, - text_properties: List[_TextProperties], - text_anchor: Position, - ) -> np.ndarray: - anchor_coordinates = detections.get_anchors_coordinates( - anchor=self.text_anchor - ).astype(int) - - xyxy = [] - for idx, center_coords in enumerate(anchor_coordinates): text_background_xyxy = resolve_text_background_xyxy( center_coordinates=tuple(center_coords), - text_wh=( - text_properties[idx].width_padded, - text_properties[idx].height_padded, - ), - position=text_anchor, + text_wh=(width_padded, height_padded), + position=self.text_anchor, ) - xyxy.append(text_background_xyxy) - return np.array(xyxy) + label_properties.append([*text_background_xyxy, text_left, text_top]) + + return np.array(label_properties).reshape(-1, 6) @staticmethod def _get_labels_text( @@ -1540,18 +1484,23 @@ def _get_labels_text( def _draw_labels( self, draw, - xyxy: np.ndarray, - text_properties: List[_TextProperties], + labels: List[str], + label_properties: np.ndarray, detections: Detections, custom_color_lookup: Optional[np.ndarray], ) -> None: + assert len(labels) == len(label_properties) == len(detections), ( + f"Number of label properties ({len(label_properties)}), " + f"labels ({len(labels)}) and detections ({len(detections)}) " + "do not match." + ) color_lookup = ( custom_color_lookup if custom_color_lookup is not None else self.color_lookup ) - for idx, box_xyxy in enumerate(xyxy): + for idx, label_property in enumerate(label_properties): background_color = resolve_color( color=self.color, detections=detections, @@ -1565,12 +1514,11 @@ def _draw_labels( color_lookup=color_lookup, ) - label_x_position = ( - box_xyxy[0] + self.text_padding - text_properties[idx].text_left - ) - label_y_position = ( - box_xyxy[1] + self.text_padding - text_properties[idx].text_top - ) + box_xyxy = label_property[:4] + text_left = label_property[4] + text_top = label_property[5] + label_x_position = box_xyxy[0] + self.text_padding - text_left + label_y_position = box_xyxy[1] + self.text_padding - text_top draw.rounded_rectangle( tuple(box_xyxy), @@ -1580,7 +1528,7 @@ def _draw_labels( ) draw.text( xy=(label_x_position, label_y_position), - text=text_properties[idx].text, + text=labels[idx], font=self.font, fill=text_color.as_rgb(), ) From 79b45a4ea8afbf8456a3df58ce9dbbde1e885b6c Mon Sep 17 00:00:00 2001 From: LinasKo Date: Mon, 11 Nov 2024 15:03:52 +0200 Subject: [PATCH 11/16] Label, RichLabel VertexLabel annotators: rename 'smart_positions' arg to 'smart_position' --- supervision/annotators/core.py | 16 ++++++++-------- supervision/keypoint/annotators.py | 8 ++++---- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index f8ab90827..b47397143 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1056,7 +1056,7 @@ def __init__( text_position: Position = Position.TOP_LEFT, color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, - smart_positions: bool = False, + smart_position: bool = False, ): """ Args: @@ -1073,7 +1073,7 @@ def __init__( Options are `INDEX`, `CLASS`, `TRACK`. border_radius (int): The radius to apply round edges. If the selected value is higher than the lower dimension, width or height, is clipped. - smart_positions (bool): Spread out the labels to avoid overlapping. + smart_position (bool): Spread out the labels to avoid overlapping. """ self.border_radius: int = border_radius self.color: Union[Color, ColorPalette] = color @@ -1083,7 +1083,7 @@ def __init__( self.text_padding: int = text_padding self.text_anchor: Position = text_position self.color_lookup: ColorLookup = color_lookup - self.smart_positions = smart_positions + self.smart_position = smart_position @ensure_cv2_image_for_annotation def annotate( @@ -1141,7 +1141,7 @@ def annotate( label_properties = self._get_label_properties(detections, labels) xyxy = label_properties[:, :4] - if self.smart_positions: + if self.smart_position: xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy @@ -1331,7 +1331,7 @@ def __init__( text_position: Position = Position.TOP_LEFT, color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, - smart_positions: bool = False, + smart_position: bool = False, ): """ Args: @@ -1348,7 +1348,7 @@ def __init__( Options are `INDEX`, `CLASS`, `TRACK`. border_radius (int): The radius to apply round edges. If the selected value is higher than the lower dimension, width or height, is clipped. - smart_positions (bool): Spread out the labels to avoid overlapping. + smart_position (bool): Spread out the labels to avoid overlapping. """ self.color = color self.text_color = text_color @@ -1356,7 +1356,7 @@ def __init__( self.text_anchor = text_position self.color_lookup = color_lookup self.border_radius = border_radius - self.smart_positions = smart_positions + self.smart_position = smart_position self.font = self._load_font(font_size, font_path) @ensure_pil_image_for_annotation @@ -1414,7 +1414,7 @@ def annotate( label_properties = self._get_label_properties(draw, detections, labels) xyxy = label_properties[:, :4] - if self.smart_positions: + if self.smart_position: xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy diff --git a/supervision/keypoint/annotators.py b/supervision/keypoint/annotators.py index fc450a837..7537b264a 100644 --- a/supervision/keypoint/annotators.py +++ b/supervision/keypoint/annotators.py @@ -202,7 +202,7 @@ def __init__( text_thickness: int = 1, text_padding: int = 10, border_radius: int = 0, - smart_positions: bool = False, + smart_position: bool = False, ): """ Args: @@ -217,7 +217,7 @@ def __init__( text_padding (int): The padding around the text. border_radius (int): The radius of the rounded corners of the boxes. Set to a high value to produce circles. - smart_positions (bool): Spread out the labels to avoid overlap. + smart_position (bool): Spread out the labels to avoid overlap. """ self.border_radius: int = border_radius self.color: Union[Color, List[Color]] = color @@ -225,7 +225,7 @@ def __init__( self.text_scale: float = text_scale self.text_thickness: int = text_thickness self.text_padding: int = text_padding - self.smart_positions = smart_positions + self.smart_position = smart_position def annotate( self, @@ -362,7 +362,7 @@ def annotate( ) xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding) - if self.smart_positions: + if self.smart_position: xyxy_padded = spread_out_boxes(xyxy_padded) xyxy = pad_boxes(xyxy=xyxy_padded, px=-self.text_padding) From 3532130b2da32f27b4ba945f290b126968267bd1 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Mon, 11 Nov 2024 15:05:57 +0200 Subject: [PATCH 12/16] Label annotators: `xyxy = label_properties[:, :4]` moved inside `if self.smart_position:` --- supervision/annotators/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index b47397143..cad8cafe8 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1139,9 +1139,9 @@ def annotate( labels = self._get_labels_text(detections, labels) label_properties = self._get_label_properties(detections, labels) - xyxy = label_properties[:, :4] if self.smart_position: + xyxy = label_properties[:, :4] xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy @@ -1412,9 +1412,9 @@ def annotate( draw = ImageDraw.Draw(scene) labels = self._get_labels_text(detections, labels) label_properties = self._get_label_properties(draw, detections, labels) - xyxy = label_properties[:, :4] if self.smart_position: + xyxy = label_properties[:, :4] xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy From 10652efbc42d8ec525089f00024df8a5db6c1d7f Mon Sep 17 00:00:00 2001 From: LinasKo Date: Mon, 11 Nov 2024 15:07:22 +0200 Subject: [PATCH 13/16] Label annotators: remove unused functions `get_box_intersection`, `get_unit_vector` --- supervision/detection/utils.py | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 246d3c607..841879446 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1041,26 +1041,6 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: return np.cross(vector_at_zero, anchors - vector_start) -def get_box_intersection( - xyxy_1: np.ndarray, xyxy_2: np.ndarray -) -> Optional[np.ndarray]: - overlap_xmin = max(xyxy_1[0], xyxy_2[0]) - overlap_ymin = max(xyxy_1[1], xyxy_2[1]) - overlap_xmax = min(xyxy_1[2], xyxy_2[2]) - overlap_ymax = min(xyxy_1[3], xyxy_2[3]) - - if overlap_xmin < overlap_xmax and overlap_ymin < overlap_ymax: - return np.array([overlap_xmin, overlap_ymin, overlap_xmax, overlap_ymax]) - else: - return None - - -def get_unit_vector(xy_1: np.ndarray, xy_2: np.ndarray) -> np.ndarray: - direction = xy_2 - xy_1 - magnitude = np.linalg.norm(direction) - return direction / magnitude if magnitude > 0 else np.zeros(2) - - def spread_out_boxes( xyxy: np.ndarray, max_iterations: int = 100, From c12c562e2958da01a20ab175acaf9fc6d978f31a Mon Sep 17 00:00:00 2001 From: LinasKo Date: Mon, 11 Nov 2024 15:15:30 +0200 Subject: [PATCH 14/16] Label annotators: add docstring to `_get_label_properties` --- supervision/annotators/core.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index cad8cafe8..02ab47d6f 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1168,6 +1168,13 @@ def _get_label_properties( detections: Detections, labels: List[str], ) -> np.ndarray: + """ + Calculate the numerical properties required to draw the labels on the image. + + Returns: + (np.ndarray): An array of label properties, containing columns: + `min_x`, `min_y`, `max_x`, `max_y`, `padded_text_height`. + """ label_properties = [] anchors_coordinates = detections.get_anchors_coordinates( anchor=self.text_anchor @@ -1439,6 +1446,15 @@ def _validate_labels(self, labels: Optional[List[str]], detections: Detections): def _get_label_properties( self, draw, detections: Detections, labels: List[str] ) -> np.ndarray: + """ + Calculate the numerical properties required to draw the labels on the image. + + Returns: + (np.ndarray): An array of label properties, containing columns: + `min_x`, `min_y`, `max_x`, `max_y`, `text_left_coordinate`, + `text_top_coordinate`. The first 4 values are already padded + with `text_padding`. + """ label_properties = [] anchor_coordinates = detections.get_anchors_coordinates( From 1d9f0b1994b4448de9b793dc4115a322fb9cced2 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Mon, 11 Nov 2024 19:47:58 +0200 Subject: [PATCH 15/16] SmartLabels: Labels move proportianlly to IoU --- supervision/detection/utils.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 841879446..2e486e48c 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1044,7 +1044,6 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: def spread_out_boxes( xyxy: np.ndarray, max_iterations: int = 100, - force_multiplier: float = 0.03, ) -> np.ndarray: """ Spread out boxes that overlap with each other. @@ -1052,8 +1051,6 @@ def spread_out_boxes( Args: xyxy: Numpy array of shape (N, 4) where N is the number of boxes. max_iterations: Maximum number of iterations to run the algorithm for. - force_multiplier: Multiplier to scale the force vectors by. Similar to - learning rate in gradient descent. """ if len(xyxy) == 0: return xyxy @@ -1076,14 +1073,26 @@ def spread_out_boxes( delta_centers *= overlap_mask[:, :, np.newaxis] # Nx2 - force_vectors = np.sum(delta_centers, axis=1) - force_vectors *= force_multiplier + delta_sum = np.sum(delta_centers, axis=1) + delta_magnitude = np.linalg.norm(delta_sum, axis=1, keepdims=True) + direction_vectors = np.divide( + delta_sum, + delta_magnitude, + out=np.zeros_like(delta_sum), + where=delta_magnitude != 0, + ) + + force_vectors = np.sum(iou, axis=1) + force_vectors = force_vectors[:, np.newaxis] * direction_vectors + + force_vectors *= 10 force_vectors[(force_vectors > 0) & (force_vectors < 1)] = 1 force_vectors[(force_vectors < 0) & (force_vectors > -1)] = -1 - # Reduce motion along primary axis - primary_axis = np.argmax(np.abs(force_vectors), axis=1) - force_vectors[np.arange(len(force_vectors)), primary_axis] /= 2 + # Move along main axis only. + force_vectors[ + np.arange(len(force_vectors)), np.argmin(force_vectors, axis=1) + ] = 0 force_vectors = force_vectors.astype(int) From 3692fb65f9c998e0e8fb74858e191a9470f27556 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Mon, 11 Nov 2024 20:07:51 +0200 Subject: [PATCH 16/16] Smart Lables: Increase minimum step size to 2 --- supervision/detection/utils.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 2e486e48c..c6c63286d 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1086,13 +1086,8 @@ def spread_out_boxes( force_vectors = force_vectors[:, np.newaxis] * direction_vectors force_vectors *= 10 - force_vectors[(force_vectors > 0) & (force_vectors < 1)] = 1 - force_vectors[(force_vectors < 0) & (force_vectors > -1)] = -1 - - # Move along main axis only. - force_vectors[ - np.arange(len(force_vectors)), np.argmin(force_vectors, axis=1) - ] = 0 + force_vectors[(force_vectors > 0) & (force_vectors < 2)] = 2 + force_vectors[(force_vectors < 0) & (force_vectors > -2)] = -2 force_vectors = force_vectors.astype(int)