From f2e900171ee09b2ff6024896e69e5394a8806007 Mon Sep 17 00:00:00 2001 From: hidara2000 Date: Tue, 15 Apr 2025 18:09:08 +1000 Subject: [PATCH 01/15] feat(annotators): enhance label annotators with frame boundary adjustments and new base class - ensures labels are within frame - May have a few overlaps at edges,in very busy frames, when smart_pos is enabled. but running smart_pos first yields better results --- supervision/annotators/core.py | 478 ++++++++++++++++++--------------- 1 file changed, 262 insertions(+), 216 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 1b5b9bee..c145229c 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,6 +1,6 @@ from functools import lru_cache from math import sqrt -from typing import List, Optional, Tuple, Union +from typing import Callable, List, Optional, Tuple, Union import cv2 import numpy as np @@ -40,6 +40,199 @@ from supervision.utils.image import ( CV2_FONT = cv2.FONT_HERSHEY_SIMPLEX +class _BaseLabelAnnotator(BaseAnnotator): + """ + Base class for annotators that add labels to detections. + + Attributes: + color (Union[Color, ColorPalette]): The color to use for the label background. + text_color (Union[Color, ColorPalette]): The color to use for the label text. + text_padding (int): The padding around the label text, in pixels. + text_anchor (Position): The position of the text relative to the detection + bounding box. + color_lookup (ColorLookup): The method used to determine the color of the label. + border_radius (int): The radius of the label background corners, in pixels. + smart_position (bool): Whether to intelligently adjust the label position to + avoid overlapping with other elements. + ensure_in_frame (bool): Whether to ensure the label stays within the frame + boundaries. + """ + + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.DEFAULT, + text_color: Union[Color, ColorPalette] = Color.WHITE, + text_padding: int = 10, + text_position: Position = Position.TOP_LEFT, + color_lookup: ColorLookup = ColorLookup.CLASS, + border_radius: int = 0, + smart_position: bool = False, + ensure_in_frame: bool = True, # New parameter + ): + """ + Initializes the _BaseLabelAnnotator. + + Args: + color (Union[Color, ColorPalette], optional): The color to use for the label + background. Defaults to ColorPalette.DEFAULT. + text_color (Union[Color, ColorPalette], optional): The color to use for the + label text. Defaults to Color.WHITE. + text_padding (int, optional): The padding around the label text, in pixels. + Defaults to 10. + text_position (Position, optional): The position of the text relative to the + detection bounding box. Defaults to Position.TOP_LEFT. + color_lookup (ColorLookup, optional): The method used to determine the color + of the label. Defaults to ColorLookup.CLASS. + border_radius (int, optional): The radius of the label background corners, + in pixels. Defaults to 0. + smart_position (bool, optional): Whether to intelligently adjust the label + position to avoid overlapping with other elements. + Defaults to False. + ensure_in_frame (bool, optional): Whether to ensure the label stays within + the frame boundaries. Defaults to True. + """ + self.color: Union[Color, ColorPalette] = color + self.text_color: Union[Color, ColorPalette] = text_color + self.text_padding: int = text_padding + self.text_anchor: Position = text_position + self.color_lookup: ColorLookup = color_lookup + self.border_radius: int = border_radius + self.smart_position = smart_position + self.ensure_in_frame = ensure_in_frame # Store the new parameter + + def _validate_labels(self, labels: Optional[List[str]], detections: Detections): + """ + Validates that the number of provided labels matches the number of detections. + + Args: + labels (Optional[List[str]]): A list of labels, one for each detection. Can + be None. + detections (Detections): The detections to be labeled. + + Raises: + ValueError: If `labels` is not None and its length does not match 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 " + f"number of detections ({len(detections)}). Each detection " + f"should have exactly 1 label." + ) + + @staticmethod + def _get_labels_text( + detections: Detections, custom_labels: Optional[List[str]] + ) -> List[str]: + """ + Retrieves the text labels for the detections. + + If `custom_labels` are provided, they are used. Otherwise, the labels are + extracted from the `detections` object, prioritizing the 'class_name' field, + then the `class_id`, and finally using the detection index as a string. + + Args: + detections (Detections): The detections to get labels for. + custom_labels (Optional[List[str]]): An optional list of custom labels. + + Returns: + List[str]: A list of text labels for each detection. + """ + 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 _adjust_labels_in_frame( + self, + frame_width: int, + frame_height: int, + labels: List[str], + label_properties: np.ndarray, + get_text_width_fn: Callable[[str], int], + ) -> np.ndarray: + """ + Adjusts the position of labels to ensure they stay within the frame boundaries. + + Args: + frame_width (int): The width of the frame. + frame_height (int): The height of the frame. + labels (List[str]): The list of text labels. + label_properties (np.ndarray): An array of label properties, where each row + contains [x1, y1, x2, y2, text_height, ...]. + get_text_width_fn (Callable[[str], int]): A function that takes a label + string and returns its width in pixels. + + Returns: + np.ndarray: The adjusted label properties. + """ + adjusted_properties = label_properties.copy() + for i in range(len(labels)): + x1, y1, x2, y2, text_height = adjusted_properties[i][ + :5 + ] # Handle different property lengths + label_h = text_height + 2 * self.text_padding + + # Adjust x-coordinate + if x1 < 0: + adjusted_properties[i, 0] -= x1 + adjusted_properties[i, 2] -= x1 + elif x2 > frame_width: + adjusted_properties[i, 0] -= x2 - frame_width + adjusted_properties[i, 2] -= x2 - frame_width + + # Adjust y-coordinate + if y1 < 0: + adjusted_properties[i, 1] -= y1 + adjusted_properties[i, 3] -= y1 + elif y2 > frame_height: + adjusted_properties[i, 1] -= y2 - frame_height + adjusted_properties[i, 3] -= y2 - frame_height + # Optionally, if the label is below the box, try to move it above + anchor_y = self._get_anchor_y_for_adjustment( + adjusted_properties[i, 1:3], self.text_anchor + ) + if anchor_y - label_h >= 0: + adjusted_properties[i, 1] -= label_h + adjusted_properties[i, 3] -= label_h + + return adjusted_properties + + @staticmethod + def _get_anchor_y_for_adjustment(bbox_y: np.ndarray, anchor: Position) -> float: + """ + Calculates the anchor y-coordinate for label adjustment based on the text anchor + position. + + Args: + bbox_y (np.ndarray): An array containing the y1 and y2 coordinates of the + bounding box. + anchor (Position): The desired text anchor position. + + Returns: + float: The anchor y-coordinate. + """ + y1, y2 = bbox_y + if anchor in [Position.TOP_LEFT, Position.TOP_CENTER, Position.TOP_RIGHT]: + return y1 + elif anchor in [ + Position.BOTTOM_LEFT, + Position.BOTTOM_CENTER, + Position.BOTTOM_RIGHT, + ]: + return y2 + else: # CENTER, CENTER_LEFT, CENTER_RIGHT + return (y1 + y2) / 2 + + class BoxAnnotator(BaseAnnotator): """ A class for drawing bounding boxes on an image using provided detections. @@ -960,7 +1153,7 @@ class DotAnnotator(BaseAnnotator): return scene -class LabelAnnotator(BaseAnnotator): +class LabelAnnotator(_BaseLabelAnnotator): """ A class for annotating labels on an image using provided detections. """ @@ -976,94 +1169,53 @@ class LabelAnnotator(BaseAnnotator): color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, smart_position: bool = False, + ensure_in_frame: bool = False, # Inherited ): - """ - Args: - color (Union[Color, ColorPalette]): The color or color palette to use for - annotating the text background. - text_color (Union[Color, ColorPalette]): The color or color palette to use - for the text. - text_scale (float): Font scale for the text. - text_thickness (int): Thickness of the text characters. - text_padding (int): Padding around the text within its background box. - text_position (Position): Position of the text relative to the detection. - Possible values are defined in the `Position` enum. - color_lookup (ColorLookup): Strategy for mapping colors to annotations. - 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_position (bool): Spread out the labels to avoid overlapping. - """ - self.border_radius: int = border_radius - self.color: Union[Color, ColorPalette] = color - self.text_color: Union[Color, ColorPalette] = text_color self.text_scale: float = text_scale self.text_thickness: int = text_thickness - self.text_padding: int = text_padding - self.text_anchor: Position = text_position - self.color_lookup: ColorLookup = color_lookup - self.smart_position = smart_position + super().__init__( + color=color, + text_color=text_color, + text_padding=text_padding, + text_position=text_position, + color_lookup=color_lookup, + border_radius=border_radius, + smart_position=smart_position, + ensure_in_frame=ensure_in_frame, + ) @ensure_cv2_image_for_annotation def annotate( self, - scene: ImageType, + scene: np.ndarray, # Ensure scene is initially a NumPy array here detections: Detections, labels: Optional[List[str]] = None, custom_color_lookup: Optional[np.ndarray] = None, - ) -> ImageType: - """ - Annotates the given scene with labels based on the provided detections. - - Args: - scene (ImageType): The image where labels will be drawn. - `ImageType` is a flexible type, accepting either `numpy.ndarray` - or `PIL.Image.Image`. - detections (Detections): Object detections to annotate. - labels (Optional[List[str]]): Custom labels for each detection. - custom_color_lookup (Optional[np.ndarray]): Custom color lookup array. - Allows to override the default color mapping strategy. - - Returns: - The annotated image, matching the type of `scene` (`numpy.ndarray` - or `PIL.Image.Image`) - - Example: - ```python - import supervision as sv - - image = ... - detections = sv.Detections(...) - - labels = [ - f"{class_name} {confidence:.2f}" - for class_name, confidence - in zip(detections['class_name'], detections.confidence) - ] - - label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER) - annotated_frame = label_annotator.annotate( - scene=image.copy(), - detections=detections, - labels=labels - ) - ``` - - ![label-annotator-example](https://media.roboflow.com/ - supervision-annotator-examples/label-annotator-example-purple.png) - """ - + ) -> np.ndarray: assert isinstance(scene, np.ndarray) self._validate_labels(labels, detections) labels = self._get_labels_text(detections, labels) - label_properties = self._get_label_properties(detections, labels) + label_properties = self._get_label_properties( + scene, detections, labels + ) # Pass scene if self.smart_position: xyxy = label_properties[:, :4] xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy + if self.ensure_in_frame: + label_properties = self._adjust_labels_in_frame( + scene.shape[1], + scene.shape[0], + labels, + label_properties, + lambda text: cv2.getTextSize( + text, CV2_FONT, self.text_scale, self.text_thickness + )[0][0], + ) + self._draw_labels( scene=scene, labels=labels, @@ -1074,26 +1226,12 @@ class LabelAnnotator(BaseAnnotator): 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"The number of labels ({len(labels)}) does not match the " - f"number of detections ({len(detections)}). Each detection " - f"should have exactly 1 label." - ) - def _get_label_properties( self, + scene: np.ndarray, # Receive scene 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 @@ -1116,32 +1254,12 @@ class LabelAnnotator(BaseAnnotator): position=self.text_anchor, ) - label_properties.append( - [ - *text_background_xyxy, - text_h, - ] - ) - + label_properties.append([ + *text_background_xyxy, + text_h, + ]) return np.array(label_properties).reshape(-1, 5) - @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, scene: np.ndarray, @@ -1176,7 +1294,7 @@ class LabelAnnotator(BaseAnnotator): color_lookup=color_lookup, ) - box_xyxy = label_property[:4] + box_xyxy = label_property[:4].astype(int) text_height_padded = label_property[4] self.draw_rounded_rectangle( scene=scene, @@ -1186,7 +1304,7 @@ class LabelAnnotator(BaseAnnotator): ) text_x = box_xyxy[0] + self.text_padding - text_y = box_xyxy[1] + self.text_padding + text_height_padded + text_y = box_xyxy[1] + self.text_padding + int(text_height_padded) cv2.putText( img=scene, text=labels[idx], @@ -1241,7 +1359,7 @@ class LabelAnnotator(BaseAnnotator): return scene -class RichLabelAnnotator(BaseAnnotator): +class RichLabelAnnotator(_BaseLabelAnnotator): """ A class for annotating labels on an image using provided detections, with support for Unicode characters by using a custom font. @@ -1258,80 +1376,30 @@ class RichLabelAnnotator(BaseAnnotator): color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, smart_position: bool = False, + ensure_in_frame: bool = False, # Inherited ): - """ - Args: - color (Union[Color, ColorPalette]): The color or color palette to use for - annotating the text background. - text_color (Union[Color, ColorPalette]): The color to use for the text. - font_path (Optional[str]): Path to the font file (e.g., ".ttf" or ".otf") - to use for rendering text. If `None`, the default PIL font will be used. - font_size (int): Font size for the text. - text_padding (int): Padding around the text within its background box. - text_position (Position): Position of the text relative to the detection. - Possible values are defined in the `Position` enum. - color_lookup (ColorLookup): Strategy for mapping colors to annotations. - 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_position (bool): Spread out the labels to avoid overlapping. - """ - self.color = color - self.text_color = text_color - self.text_padding = text_padding - self.text_anchor = text_position - self.color_lookup = color_lookup - self.border_radius = border_radius - self.smart_position = smart_position + self.font_path = font_path + self.font_size = font_size self.font = self._load_font(font_size, font_path) + super().__init__( + color=color, + text_color=text_color, + text_padding=text_padding, + text_position=text_position, + color_lookup=color_lookup, + border_radius=border_radius, + smart_position=smart_position, + ensure_in_frame=ensure_in_frame, + ) @ensure_pil_image_for_annotation def annotate( self, - scene: ImageType, + scene: Image.Image, detections: Detections, labels: Optional[List[str]] = None, custom_color_lookup: Optional[np.ndarray] = None, - ) -> ImageType: - """ - Annotates the given scene with labels based on the provided - detections, with support for Unicode characters. - - Args: - scene (ImageType): The image where labels will be drawn. - `ImageType` is a flexible type, accepting either `numpy.ndarray` - or `PIL.Image.Image`. - detections (Detections): Object detections to annotate. - labels (Optional[List[str]]): Custom labels for each detection. - custom_color_lookup (Optional[np.ndarray]): Custom color lookup array. - Allows to override the default color mapping strategy. - - Returns: - The annotated image, matching the type of `scene` (`numpy.ndarray` - or `PIL.Image.Image`) - - Example: - ```python - import supervision as sv - - image = ... - detections = sv.Detections(...) - - labels = [ - f"{class_name} {confidence:.2f}" - for class_name, confidence - in zip(detections['class_name'], detections.confidence) - ] - - rich_label_annotator = sv.RichLabelAnnotator(font_path="path/to/font.ttf") - annotated_frame = label_annotator.annotate( - scene=image.copy(), - detections=detections, - labels=labels - ) - ``` - - """ + ) -> Image.Image: assert isinstance(scene, Image.Image) self._validate_labels(labels, detections) @@ -1344,6 +1412,18 @@ class RichLabelAnnotator(BaseAnnotator): xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy + if self.ensure_in_frame: + label_properties = self._adjust_labels_in_frame( + scene.width, + scene.height, + labels, + label_properties, + lambda text: int( + draw.textbbox((0, 0), text, font=self.font)[2] + - draw.textbbox((0, 0), text, font=self.font)[0] + ), + ) + self._draw_labels( draw=draw, labels=labels, @@ -1354,26 +1434,9 @@ class RichLabelAnnotator(BaseAnnotator): 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"The number of labels ({len(labels)}) does not match the " - f"number of detections ({len(detections)}). Each detection " - f"should have exactly 1 label." - ) - def _get_label_properties( - self, draw, detections: Detections, labels: List[str] + self, draw: ImageDraw.ImageDraw, 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( @@ -1386,8 +1449,8 @@ class RichLabelAnnotator(BaseAnnotator): ) 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 + width_padded = int(text_width + 2 * self.text_padding) + height_padded = int(text_height + 2 * self.text_padding) text_background_xyxy = resolve_text_background_xyxy( center_coordinates=tuple(center_coords), @@ -1399,26 +1462,9 @@ class RichLabelAnnotator(BaseAnnotator): return np.array(label_properties).reshape(-1, 6) - @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, + draw: ImageDraw.ImageDraw, labels: List[str], label_properties: np.ndarray, detections: Detections, @@ -1449,7 +1495,7 @@ class RichLabelAnnotator(BaseAnnotator): color_lookup=color_lookup, ) - box_xyxy = label_property[:4] + box_xyxy = label_property[:4].astype(int) text_left = label_property[4] text_top = label_property[5] label_x_position = box_xyxy[0] + self.text_padding - text_left From ca5596ceb2d87346f0661fcd8eee593a7ff22be9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 15 Apr 2025 08:30:24 +0000 Subject: [PATCH 02/15] =?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 | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index c145229c..d35a1eba 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1254,10 +1254,12 @@ class LabelAnnotator(_BaseLabelAnnotator): position=self.text_anchor, ) - label_properties.append([ - *text_background_xyxy, - text_h, - ]) + label_properties.append( + [ + *text_background_xyxy, + text_h, + ] + ) return np.array(label_properties).reshape(-1, 5) def _draw_labels( From a1e43f8e933052896d5eb7987c30bdd3b626fd3f Mon Sep 17 00:00:00 2001 From: hidara2000 <15170494+hidara2000@users.noreply.github.com> Date: Tue, 15 Apr 2025 21:58:06 +1000 Subject: [PATCH 03/15] Update supervision/annotators/core.py Co-authored-by: Onuralp SEZER --- supervision/annotators/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index d35a1eba..4f261888 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1187,7 +1187,7 @@ class LabelAnnotator(_BaseLabelAnnotator): @ensure_cv2_image_for_annotation def annotate( self, - scene: np.ndarray, # Ensure scene is initially a NumPy array here + scene: ImageType, # Ensure scene is initially a NumPy array here detections: Detections, labels: Optional[List[str]] = None, custom_color_lookup: Optional[np.ndarray] = None, From a5c833b80f163dfceab7f304c9e4f435139ec53d Mon Sep 17 00:00:00 2001 From: hidara2000 <15170494+hidara2000@users.noreply.github.com> Date: Tue, 15 Apr 2025 21:58:20 +1000 Subject: [PATCH 04/15] Update supervision/annotators/core.py Co-authored-by: Onuralp SEZER --- supervision/annotators/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 4f261888..99c23ed7 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1397,7 +1397,7 @@ class RichLabelAnnotator(_BaseLabelAnnotator): @ensure_pil_image_for_annotation def annotate( self, - scene: Image.Image, + scene: ImageType detections: Detections, labels: Optional[List[str]] = None, custom_color_lookup: Optional[np.ndarray] = None, From 474c8bc478c7a09b009072f0db9fff4c3a4508c6 Mon Sep 17 00:00:00 2001 From: hidara2000 <15170494+hidara2000@users.noreply.github.com> Date: Tue, 15 Apr 2025 21:58:31 +1000 Subject: [PATCH 05/15] Update supervision/annotators/core.py Co-authored-by: Onuralp SEZER --- supervision/annotators/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 99c23ed7..0fee71d1 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1401,7 +1401,7 @@ class RichLabelAnnotator(_BaseLabelAnnotator): detections: Detections, labels: Optional[List[str]] = None, custom_color_lookup: Optional[np.ndarray] = None, - ) -> Image.Image: + ) -> ImageType: assert isinstance(scene, Image.Image) self._validate_labels(labels, detections) From c4dce2752d2479404abfed241196ad1e5e0a2117 Mon Sep 17 00:00:00 2001 From: hidara2000 Date: Tue, 15 Apr 2025 22:09:51 +1000 Subject: [PATCH 06/15] fix(annotators): add missing comma in annotate method signature --- supervision/annotators/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 0fee71d1..f6e71e2e 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1397,7 +1397,7 @@ class RichLabelAnnotator(_BaseLabelAnnotator): @ensure_pil_image_for_annotation def annotate( self, - scene: ImageType + scene: ImageType, detections: Detections, labels: Optional[List[str]] = None, custom_color_lookup: Optional[np.ndarray] = None, From 0539ec9d1f63e1b87f8e5e370ba9949f00747cdf Mon Sep 17 00:00:00 2001 From: hidara2000 Date: Thu, 17 Apr 2025 10:48:02 +1000 Subject: [PATCH 07/15] feat(annotators): add max_line_length parameter and text wrapping functionality --- supervision/annotators/core.py | 327 +++++++++++++++++++++++++-------- supervision/detection/utils.py | 65 +++++-- 2 files changed, 302 insertions(+), 90 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index f6e71e2e..254c55e0 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -54,8 +54,8 @@ class _BaseLabelAnnotator(BaseAnnotator): border_radius (int): The radius of the label background corners, in pixels. smart_position (bool): Whether to intelligently adjust the label position to avoid overlapping with other elements. - ensure_in_frame (bool): Whether to ensure the label stays within the frame - boundaries. + max_line_length (Optional[int]): Maximum number of characters per line before + wrapping the text. None means no wrapping. """ def __init__( @@ -67,7 +67,7 @@ class _BaseLabelAnnotator(BaseAnnotator): color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, smart_position: bool = False, - ensure_in_frame: bool = True, # New parameter + max_line_length: Optional[int] = None, ): """ Initializes the _BaseLabelAnnotator. @@ -88,8 +88,9 @@ class _BaseLabelAnnotator(BaseAnnotator): smart_position (bool, optional): Whether to intelligently adjust the label position to avoid overlapping with other elements. Defaults to False. - ensure_in_frame (bool, optional): Whether to ensure the label stays within - the frame boundaries. Defaults to True. + max_line_length (Optional[int], optional): Maximum number of characters per + line before wrapping the text. None means no wrapping. + Defaults to None. """ self.color: Union[Color, ColorPalette] = color self.text_color: Union[Color, ColorPalette] = text_color @@ -98,7 +99,7 @@ class _BaseLabelAnnotator(BaseAnnotator): self.color_lookup: ColorLookup = color_lookup self.border_radius: int = border_radius self.smart_position = smart_position - self.ensure_in_frame = ensure_in_frame # Store the new parameter + self.max_line_length: Optional[int] = max_line_length def _validate_labels(self, labels: Optional[List[str]], detections: Detections): """ @@ -151,6 +152,67 @@ class _BaseLabelAnnotator(BaseAnnotator): labels.append(str(idx)) return labels + def _ensure_box_in_frame( + self, + box: np.ndarray, + frame_width: int, + frame_height: int, + check_flip_label: bool = False, + text_anchor: Optional[Position] = None, + ) -> np.ndarray: + """ + Adjusts a box to ensure it stays within frame boundaries. + + Args: + box: Box coordinates [x1, y1, x2, y2] + frame_width: Width of the frame + frame_height: Height of the frame + check_flip_label: Whether to check if label should be flipped + text_anchor: Text anchor position for flip check + + Returns: + Adjusted box coordinates + """ + x1, y1, x2, y2 = box + + # Adjust x-coordinate to stay within frame + if x1 < 0: + shift = -x1 + x1 += shift + x2 += shift + elif x2 > frame_width: + shift = frame_width - x2 + x1 += shift + x2 += shift + + # Adjust y-coordinate to stay within frame + if y1 < 0: + shift = -y1 + y1 += shift + y2 += shift + elif y2 > frame_height: + shift = frame_height - y2 + y1 += shift + y2 += shift + + # Check if label should be flipped to above the box + if check_flip_label and text_anchor is not None: + box_height = y2 - y1 + + # Check anchor position to see if we can flip it + anchor_y = self._get_anchor_y_for_adjustment( + np.array([y1, y2]), text_anchor + ) + + # If we're at the bottom, try moving to the top + if anchor_y >= y2 - 5: # Near bottom edge + # Check if there's room at the top + if y1 - box_height >= 0: + y2 = y1 + y1 = y2 - box_height + + return np.array([x1, y1, x2, y2]) + def _adjust_labels_in_frame( self, frame_width: int, @@ -175,34 +237,33 @@ class _BaseLabelAnnotator(BaseAnnotator): np.ndarray: The adjusted label properties. """ adjusted_properties = label_properties.copy() + + # First, make sure the boxes don't go outside the frame for i in range(len(labels)): - x1, y1, x2, y2, text_height = adjusted_properties[i][ - :5 - ] # Handle different property lengths - label_h = text_height + 2 * self.text_padding + # Adjust box to stay within frame + adjusted_properties[i, :4] = self._ensure_box_in_frame( + adjusted_properties[i, :4], + frame_width, + frame_height, + check_flip_label=True, + text_anchor=self.text_anchor, + ) - # Adjust x-coordinate - if x1 < 0: - adjusted_properties[i, 0] -= x1 - adjusted_properties[i, 2] -= x1 - elif x2 > frame_width: - adjusted_properties[i, 0] -= x2 - frame_width - adjusted_properties[i, 2] -= x2 - frame_width + # Apply the spread out algorithm to avoid box overlaps + if len(labels) > 1: + # Extract the box coordinates + boxes = adjusted_properties[:, :4] + # Use the spread_out_boxes function to adjust overlapping boxes + spread_boxes = spread_out_boxes(boxes) + # Update the properties with the spread out boxes + adjusted_properties[:, :4] = spread_boxes - # Adjust y-coordinate - if y1 < 0: - adjusted_properties[i, 1] -= y1 - adjusted_properties[i, 3] -= y1 - elif y2 > frame_height: - adjusted_properties[i, 1] -= y2 - frame_height - adjusted_properties[i, 3] -= y2 - frame_height - # Optionally, if the label is below the box, try to move it above - anchor_y = self._get_anchor_y_for_adjustment( - adjusted_properties[i, 1:3], self.text_anchor + # Additional check to ensure boxes are still within frame after spreading + for i in range(len(labels)): + # Adjust box to stay within frame (without flipping) + adjusted_properties[i, :4] = self._ensure_box_in_frame( + adjusted_properties[i, :4], frame_width, frame_height ) - if anchor_y - label_h >= 0: - adjusted_properties[i, 1] -= label_h - adjusted_properties[i, 3] -= label_h return adjusted_properties @@ -232,6 +293,53 @@ class _BaseLabelAnnotator(BaseAnnotator): else: # CENTER, CENTER_LEFT, CENTER_RIGHT return (y1 + y2) / 2 + def _wrap_text(self, text: str) -> List[str]: + """ + Wraps text to the specified maximum line length, respecting existing newlines. + Uses the textwrap library for robust text wrapping. + + Args: + text (str): The text to wrap. + + Returns: + List[str]: A list of text lines after wrapping. + """ + import textwrap + + if not text: + return [""] + + if self.max_line_length is None: + return text.splitlines() or [""] + + # Split the text by existing newlines first + paragraphs = text.split("\n") + all_lines = [] + + for paragraph in paragraphs: + if not paragraph: + # Keep empty lines + all_lines.append("") + continue + + # Wrap each paragraph separately + wrapped = textwrap.wrap( + paragraph, + width=self.max_line_length, + break_long_words=True, # Break words longer than max_line_length + replace_whitespace=False, # Preserve existing spaces + drop_whitespace=True, # Drop leading/trailing whitespace on wrapped lines + ) + + # Add the wrapped lines for this paragraph + if wrapped: + all_lines.extend(wrapped) + else: + # If wrap returns an empty list (e.g., for whitespace-only input) + all_lines.append("") + + return all_lines if all_lines else [""] + class BoxAnnotator(BaseAnnotator): """ @@ -1169,7 +1277,7 @@ class LabelAnnotator(_BaseLabelAnnotator): color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, smart_position: bool = False, - ensure_in_frame: bool = False, # Inherited + max_line_length: Optional[int] = None, ): self.text_scale: float = text_scale self.text_thickness: int = text_thickness @@ -1181,7 +1289,7 @@ class LabelAnnotator(_BaseLabelAnnotator): color_lookup=color_lookup, border_radius=border_radius, smart_position=smart_position, - ensure_in_frame=ensure_in_frame, + max_line_length=max_line_length, ) @ensure_cv2_image_for_annotation @@ -1205,7 +1313,6 @@ class LabelAnnotator(_BaseLabelAnnotator): xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy - if self.ensure_in_frame: label_properties = self._adjust_labels_in_frame( scene.shape[1], scene.shape[0], @@ -1238,15 +1345,29 @@ class LabelAnnotator(_BaseLabelAnnotator): ).astype(int) for label, center_coords in zip(labels, anchors_coordinates): - (text_w, text_h) = cv2.getTextSize( - text=label, - fontFace=CV2_FONT, - fontScale=self.text_scale, - thickness=self.text_thickness, - )[0] + wrapped_lines = self._wrap_text(label) + line_heights = [] + line_widths = [] - width_padded = text_w + 2 * self.text_padding - height_padded = text_h + 2 * self.text_padding + for line in wrapped_lines: + (text_w, text_h) = cv2.getTextSize( + text=line, + fontFace=CV2_FONT, + fontScale=self.text_scale, + thickness=self.text_thickness, + )[0] + line_heights.append(text_h) + line_widths.append(text_w) + + # Get the maximum width and total height + max_width = max(line_widths) if line_widths else 0 + total_height = ( + sum(line_heights) + (len(line_heights) - 1) * self.text_padding + ) + + # Add padding around all sides + width_padded = max_width + 2 * self.text_padding + height_padded = total_height + 2 * self.text_padding text_background_xyxy = resolve_text_background_xyxy( center_coordinates=tuple(center_coords), @@ -1254,12 +1375,10 @@ class LabelAnnotator(_BaseLabelAnnotator): position=self.text_anchor, ) - label_properties.append( - [ - *text_background_xyxy, - text_h, - ] - ) + label_properties.append([ + *text_background_xyxy, + total_height, + ]) return np.array(label_properties).reshape(-1, 5) def _draw_labels( @@ -1297,7 +1416,7 @@ class LabelAnnotator(_BaseLabelAnnotator): ) box_xyxy = label_property[:4].astype(int) - text_height_padded = label_property[4] + self.draw_rounded_rectangle( scene=scene, xyxy=box_xyxy, @@ -1305,18 +1424,43 @@ class LabelAnnotator(_BaseLabelAnnotator): border_radius=self.border_radius, ) - text_x = box_xyxy[0] + self.text_padding - text_y = box_xyxy[1] + self.text_padding + int(text_height_padded) - cv2.putText( - img=scene, - text=labels[idx], - org=(text_x, text_y), - fontFace=CV2_FONT, - fontScale=self.text_scale, - color=text_color.as_bgr(), - thickness=self.text_thickness, - lineType=cv2.LINE_AA, - ) + # Handle multiline text + wrapped_lines = self._wrap_text(labels[idx]) + current_y = box_xyxy[1] + self.text_padding # Start y position + + for line in wrapped_lines: + if not line: # Skip empty lines + (_, text_h) = cv2.getTextSize( + text="Tg", # Use a character with ascenders and descenders as height reference + fontFace=CV2_FONT, + fontScale=self.text_scale, + thickness=self.text_thickness, + )[0] + current_y += text_h + self.text_padding + continue + + (_, text_h) = cv2.getTextSize( + text=line, + fontFace=CV2_FONT, + fontScale=self.text_scale, + thickness=self.text_thickness, + )[0] + + text_x = box_xyxy[0] + self.text_padding + text_y = current_y + text_h # Add height to get to text baseline + + cv2.putText( + img=scene, + text=line, + org=(text_x, text_y), + fontFace=CV2_FONT, + fontScale=self.text_scale, + color=text_color.as_bgr(), + thickness=self.text_thickness, + lineType=cv2.LINE_AA, + ) + + current_y += text_h + self.text_padding # Move to next line position @staticmethod def draw_rounded_rectangle( @@ -1378,7 +1522,7 @@ class RichLabelAnnotator(_BaseLabelAnnotator): color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, smart_position: bool = False, - ensure_in_frame: bool = False, # Inherited + max_line_length: Optional[int] = None, ): self.font_path = font_path self.font_size = font_size @@ -1391,7 +1535,7 @@ class RichLabelAnnotator(_BaseLabelAnnotator): color_lookup=color_lookup, border_radius=border_radius, smart_position=smart_position, - ensure_in_frame=ensure_in_frame, + max_line_length=max_line_length, ) @ensure_pil_image_for_annotation @@ -1414,7 +1558,6 @@ class RichLabelAnnotator(_BaseLabelAnnotator): xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy - if self.ensure_in_frame: label_properties = self._adjust_labels_in_frame( scene.width, scene.height, @@ -1446,13 +1589,26 @@ class RichLabelAnnotator(_BaseLabelAnnotator): ).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 - ) - text_width = text_right - text_left - text_height = text_bottom - text_top - width_padded = int(text_width + 2 * self.text_padding) - height_padded = int(text_height + 2 * self.text_padding) + wrapped_lines = self._wrap_text(label) + + # Calculate the total text height and maximum width + max_width = 0 + total_height = 0 + + for line in wrapped_lines: + left, top, right, bottom = draw.textbbox((0, 0), line, font=self.font) + line_width = right - left + line_height = bottom - top + + max_width = max(max_width, line_width) + total_height += line_height + + # Add inter-line spacing + if len(wrapped_lines) > 1: + total_height += (len(wrapped_lines) - 1) * self.text_padding + + width_padded = int(max_width + 2 * self.text_padding) + height_padded = int(total_height + 2 * self.text_padding) text_background_xyxy = resolve_text_background_xyxy( center_coordinates=tuple(center_coords), @@ -1460,6 +1616,9 @@ class RichLabelAnnotator(_BaseLabelAnnotator): position=self.text_anchor, ) + # Get the text origin offsets + text_left, text_top, _, _ = draw.textbbox((0, 0), "Tg", font=self.font) + label_properties.append([*text_background_xyxy, text_left, text_top]) return np.array(label_properties).reshape(-1, 6) @@ -1500,21 +1659,32 @@ class RichLabelAnnotator(_BaseLabelAnnotator): box_xyxy = label_property[:4].astype(int) 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 the rounded rectangle background draw.rounded_rectangle( tuple(box_xyxy), radius=self.border_radius, fill=background_color.as_rgb(), outline=None, ) - draw.text( - xy=(label_x_position, label_y_position), - text=labels[idx], - font=self.font, - fill=text_color.as_rgb(), - ) + + # Draw each line of text + wrapped_lines = self._wrap_text(labels[idx]) + x_position = box_xyxy[0] + self.text_padding - text_left + y_position = box_xyxy[1] + self.text_padding - text_top + + for line in wrapped_lines: + draw.text( + xy=(x_position, y_position), + text=line, + font=self.font, + fill=text_color.as_rgb(), + ) + + # Move to the next line position + left, top, right, bottom = draw.textbbox((0, 0), line, font=self.font) + line_height = bottom - top + y_position += line_height + self.text_padding @staticmethod def _load_font(font_size: int, font_path: Optional[str]): @@ -1534,6 +1704,7 @@ class RichLabelAnnotator(_BaseLabelAnnotator): 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 61d0c0cb..d2191c9f 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1214,9 +1214,10 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: Returns: Array of cross products of shape (number of anchors, detections) """ - vector_at_zero = np.array( - [vector.end.x - vector.start.x, vector.end.y - vector.start.y] - ) + vector_at_zero = np.array([ + vector.end.x - vector.start.x, + vector.end.y - vector.start.y, + ]) vector_start = np.array([vector.start.x, vector.start.y]) return np.cross(vector_at_zero, anchors - vector_start) @@ -1224,6 +1225,8 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: def spread_out_boxes( xyxy: np.ndarray, max_iterations: int = 100, + force_scale: float = 10.0, + consider_size: bool = True, ) -> np.ndarray: """ Spread out boxes that overlap with each other. @@ -1231,29 +1234,50 @@ 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_scale: Scale factor for the repulsion forces. + consider_size: Whether to consider box size when calculating forces. """ if len(xyxy) == 0: return xyxy xyxy_padded = pad_boxes(xyxy, px=1) + + # Calculate box areas if we're considering size + if consider_size: + box_areas = (xyxy_padded[:, 2] - xyxy_padded[:, 0]) * ( + xyxy_padded[:, 3] - xyxy_padded[:, 1] + ) + # Calculate the size factors (normalize by mean size) + size_factors = ( + np.sqrt(box_areas) / np.mean(np.sqrt(box_areas)) + if len(box_areas) > 0 + else np.ones(len(xyxy_padded)) + ) + size_factors = np.clip(size_factors, 0.5, 2.0) # Clip to avoid extreme values + else: + size_factors = np.ones(len(xyxy_padded)) + for _ in range(max_iterations): - # NxN + # NxN matrix of IoU values iou = box_iou_batch(xyxy_padded, xyxy_padded) - np.fill_diagonal(iou, 0) + np.fill_diagonal(iou, 0) # Eliminate self-interactions + if np.all(iou == 0): - break + break # No more overlaps overlap_mask = iou > 0 - # Nx2 + # Nx2 array of box centers centers = (xyxy_padded[:, :2] + xyxy_padded[:, 2:]) / 2 - # NxNx2 + # NxNx2 array of differences between centers delta_centers = centers[:, np.newaxis, :] - centers[np.newaxis, :, :] delta_centers *= overlap_mask[:, :, np.newaxis] - # Nx2 + # Nx2 array of total delta vectors delta_sum = np.sum(delta_centers, axis=1) + + # Normalize to unit vectors delta_magnitude = np.linalg.norm(delta_sum, axis=1, keepdims=True) direction_vectors = np.divide( delta_sum, @@ -1262,15 +1286,32 @@ def spread_out_boxes( where=delta_magnitude != 0, ) + # Calculate force based on IoU values 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 < 2)] = 2 - force_vectors[(force_vectors < 0) & (force_vectors > -2)] = -2 + # Apply size-based scaling if enabled + if consider_size: + force_vectors *= size_factors[:, np.newaxis] + # Scale forces + force_vectors *= force_scale + + # Ensure minimum force for small overlaps + force_magnitudes = np.linalg.norm(force_vectors, axis=1, keepdims=True) + small_force_mask = (force_magnitudes > 0) & (force_magnitudes < 2) + if np.any(small_force_mask): + force_directions = force_vectors / np.where( + force_magnitudes > 0, force_magnitudes, 1 + ) + force_vectors = np.where( + small_force_mask, force_directions * 2, force_vectors + ) + + # Convert to integer displacements force_vectors = force_vectors.astype(int) + # Apply forces to update box positions xyxy_padded[:, [0, 1]] += force_vectors xyxy_padded[:, [2, 3]] += force_vectors From c68218a07ef3c97c73bd12843af9ef9ebb6025b1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 17 Apr 2025 00:48:21 +0000 Subject: [PATCH 08/15] =?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 | 11 ++++++----- supervision/detection/utils.py | 10 ++++++---- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 254c55e0..b78b3b17 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1375,10 +1375,12 @@ class LabelAnnotator(_BaseLabelAnnotator): position=self.text_anchor, ) - label_properties.append([ - *text_background_xyxy, - total_height, - ]) + label_properties.append( + [ + *text_background_xyxy, + total_height, + ] + ) return np.array(label_properties).reshape(-1, 5) def _draw_labels( @@ -1704,7 +1706,6 @@ class RichLabelAnnotator(_BaseLabelAnnotator): 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 d2191c9f..40275831 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1214,10 +1214,12 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: Returns: Array of cross products of shape (number of anchors, detections) """ - vector_at_zero = np.array([ - vector.end.x - vector.start.x, - vector.end.y - vector.start.y, - ]) + vector_at_zero = np.array( + [ + vector.end.x - vector.start.x, + vector.end.y - vector.start.y, + ] + ) vector_start = np.array([vector.start.x, vector.start.y]) return np.cross(vector_at_zero, anchors - vector_start) From 5c35af99d00f29a6cef05e705251c525ea771190 Mon Sep 17 00:00:00 2001 From: hidara2000 Date: Thu, 17 Apr 2025 11:02:25 +1000 Subject: [PATCH 09/15] fix long comment lines --- supervision/annotators/core.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index b78b3b17..85269821 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -326,9 +326,9 @@ class _BaseLabelAnnotator(BaseAnnotator): wrapped = textwrap.wrap( paragraph, width=self.max_line_length, - break_long_words=True, # Break words longer than max_line_length - replace_whitespace=False, # Preserve existing spaces - drop_whitespace=True, # Drop leading/trailing whitespace on wrapped lines + break_long_words=True, + replace_whitespace=False, + drop_whitespace=True, ) # Add the wrapped lines for this paragraph @@ -1375,12 +1375,10 @@ class LabelAnnotator(_BaseLabelAnnotator): position=self.text_anchor, ) - label_properties.append( - [ - *text_background_xyxy, - total_height, - ] - ) + label_properties.append([ + *text_background_xyxy, + total_height, + ]) return np.array(label_properties).reshape(-1, 5) def _draw_labels( @@ -1431,9 +1429,10 @@ class LabelAnnotator(_BaseLabelAnnotator): current_y = box_xyxy[1] + self.text_padding # Start y position for line in wrapped_lines: - if not line: # Skip empty lines + if not line: + # Use a character with ascenders and descenders as height reference (_, text_h) = cv2.getTextSize( - text="Tg", # Use a character with ascenders and descenders as height reference + text="Tg", fontFace=CV2_FONT, fontScale=self.text_scale, thickness=self.text_thickness, From 95a2ecc49b3a3a229dc8d25edb075d6772d14de4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 17 Apr 2025 01:02:42 +0000 Subject: [PATCH 10/15] =?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 | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 85269821..1fd85a71 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1375,10 +1375,12 @@ class LabelAnnotator(_BaseLabelAnnotator): position=self.text_anchor, ) - label_properties.append([ - *text_background_xyxy, - total_height, - ]) + label_properties.append( + [ + *text_background_xyxy, + total_height, + ] + ) return np.array(label_properties).reshape(-1, 5) def _draw_labels( From 29695055d298ff2be67c657040c1aefd7111d7be Mon Sep 17 00:00:00 2001 From: hidara2000 Date: Fri, 25 Apr 2025 09:32:54 +1000 Subject: [PATCH 11/15] refactor(annotators): streamline label handling and improve box snapping logic --- supervision/annotators/core.py | 280 +++++--------------------------- supervision/annotators/utils.py | 181 ++++++++++++++++++++- supervision/detection/utils.py | 68 ++------ 3 files changed, 228 insertions(+), 301 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 882b50ca..5190fa21 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,6 +1,6 @@ from functools import lru_cache from math import sqrt -from typing import Callable, List, Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import cv2 import numpy as np @@ -11,10 +11,14 @@ from supervision.annotators.base import BaseAnnotator, ImageType from supervision.annotators.utils import ( ColorLookup, Trace, + get_labels_text, resolve_color, resolve_text_background_xyxy, + snap_boxes, + validate_labels, + wrap_text, ) -from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES +from supervision.config import ORIENTED_BOX_COORDINATES from supervision.detection.core import Detections from supervision.detection.utils import ( clip_boxes, @@ -74,23 +78,21 @@ class _BaseLabelAnnotator(BaseAnnotator): Args: color (Union[Color, ColorPalette], optional): The color to use for the label - background. Defaults to ColorPalette.DEFAULT. + background. text_color (Union[Color, ColorPalette], optional): The color to use for the - label text. Defaults to Color.WHITE. + label text. text_padding (int, optional): The padding around the label text, in pixels. - Defaults to 10. text_position (Position, optional): The position of the text relative to the - detection bounding box. Defaults to Position.TOP_LEFT. + detection bounding box. color_lookup (ColorLookup, optional): The method used to determine the color - of the label. Defaults to ColorLookup.CLASS. + of the label border_radius (int, optional): The radius of the label background corners, - in pixels. Defaults to 0. + in pixels. smart_position (bool, optional): Whether to intelligently adjust the label position to avoid overlapping with other elements. - Defaults to False. max_line_length (Optional[int], optional): Maximum number of characters per line before wrapping the text. None means no wrapping. - Defaults to None. + """ self.color: Union[Color, ColorPalette] = color self.text_color: Union[Color, ColorPalette] = text_color @@ -101,137 +103,21 @@ class _BaseLabelAnnotator(BaseAnnotator): self.smart_position = smart_position self.max_line_length: Optional[int] = max_line_length - def _validate_labels(self, labels: Optional[List[str]], detections: Detections): - """ - Validates that the number of provided labels matches the number of detections. - - Args: - labels (Optional[List[str]]): A list of labels, one for each detection. Can - be None. - detections (Detections): The detections to be labeled. - - Raises: - ValueError: If `labels` is not None and its length does not match 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 " - f"number of detections ({len(detections)}). Each detection " - f"should have exactly 1 label." - ) - - @staticmethod - def _get_labels_text( - detections: Detections, custom_labels: Optional[List[str]] - ) -> List[str]: - """ - Retrieves the text labels for the detections. - - If `custom_labels` are provided, they are used. Otherwise, the labels are - extracted from the `detections` object, prioritizing the 'class_name' field, - then the `class_id`, and finally using the detection index as a string. - - Args: - detections (Detections): The detections to get labels for. - custom_labels (Optional[List[str]]): An optional list of custom labels. - - Returns: - List[str]: A list of text labels for each detection. - """ - if custom_labels is not None: - return [str(label) for label in 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 _ensure_box_in_frame( - self, - box: np.ndarray, - frame_width: int, - frame_height: int, - check_flip_label: bool = False, - text_anchor: Optional[Position] = None, - ) -> np.ndarray: - """ - Adjusts a box to ensure it stays within frame boundaries. - - Args: - box: Box coordinates [x1, y1, x2, y2] - frame_width: Width of the frame - frame_height: Height of the frame - check_flip_label: Whether to check if label should be flipped - text_anchor: Text anchor position for flip check - - Returns: - Adjusted box coordinates - """ - x1, y1, x2, y2 = box - - # Adjust x-coordinate to stay within frame - if x1 < 0: - shift = -x1 - x1 += shift - x2 += shift - elif x2 > frame_width: - shift = frame_width - x2 - x1 += shift - x2 += shift - - # Adjust y-coordinate to stay within frame - if y1 < 0: - shift = -y1 - y1 += shift - y2 += shift - elif y2 > frame_height: - shift = frame_height - y2 - y1 += shift - y2 += shift - - # Check if label should be flipped to above the box - if check_flip_label and text_anchor is not None: - box_height = y2 - y1 - - # Check anchor position to see if we can flip it - anchor_y = self._get_anchor_y_for_adjustment( - np.array([y1, y2]), text_anchor - ) - - # If we're at the bottom, try moving to the top - if anchor_y >= y2 - 5: # Near bottom edge - # Check if there's room at the top - if y1 - box_height >= 0: - y2 = y1 - y1 = y2 - box_height - - return np.array([x1, y1, x2, y2]) - def _adjust_labels_in_frame( self, - frame_width: int, - frame_height: int, + resolution_wh: Tuple[int, int], labels: List[str], label_properties: np.ndarray, - get_text_width_fn: Callable[[str], int], ) -> np.ndarray: """ Adjusts the position of labels to ensure they stay within the frame boundaries. Args: frame_width (int): The width of the frame. - frame_height (int): The height of the frame. + resolution_wh (int, int): The width and height of the frame. labels (List[str]): The list of text labels. label_properties (np.ndarray): An array of label properties, where each row contains [x1, y1, x2, y2, text_height, ...]. - get_text_width_fn (Callable[[str], int]): A function that takes a label - string and returns its width in pixels. Returns: np.ndarray: The adjusted label properties. @@ -239,15 +125,10 @@ class _BaseLabelAnnotator(BaseAnnotator): adjusted_properties = label_properties.copy() # First, make sure the boxes don't go outside the frame - for i in range(len(labels)): - # Adjust box to stay within frame - adjusted_properties[i, :4] = self._ensure_box_in_frame( - adjusted_properties[i, :4], - frame_width, - frame_height, - check_flip_label=True, - text_anchor=self.text_anchor, - ) + adjusted_properties[:, :4] = snap_boxes( + adjusted_properties[:, :4], + resolution_wh, + ) # Apply the spread out algorithm to avoid box overlaps if len(labels) > 1: @@ -259,86 +140,11 @@ class _BaseLabelAnnotator(BaseAnnotator): adjusted_properties[:, :4] = spread_boxes # Additional check to ensure boxes are still within frame after spreading - for i in range(len(labels)): - # Adjust box to stay within frame (without flipping) - adjusted_properties[i, :4] = self._ensure_box_in_frame( - adjusted_properties[i, :4], frame_width, frame_height - ) - - return adjusted_properties - - @staticmethod - def _get_anchor_y_for_adjustment(bbox_y: np.ndarray, anchor: Position) -> float: - """ - Calculates the anchor y-coordinate for label adjustment based on the text anchor - position. - - Args: - bbox_y (np.ndarray): An array containing the y1 and y2 coordinates of the - bounding box. - anchor (Position): The desired text anchor position. - - Returns: - float: The anchor y-coordinate. - """ - y1, y2 = bbox_y - if anchor in [Position.TOP_LEFT, Position.TOP_CENTER, Position.TOP_RIGHT]: - return y1 - elif anchor in [ - Position.BOTTOM_LEFT, - Position.BOTTOM_CENTER, - Position.BOTTOM_RIGHT, - ]: - return y2 - else: # CENTER, CENTER_LEFT, CENTER_RIGHT - return (y1 + y2) / 2 - - def _wrap_text(self, text: str) -> List[str]: - """ - Wraps text to the specified maximum line length, respecting existing newlines. - Uses the textwrap library for robust text wrapping. - - Args: - text (str): The text to wrap. - - Returns: - List[str]: A list of text lines after wrapping. - """ - import textwrap - - if not text: - return [""] - - if self.max_line_length is None: - return text.splitlines() or [""] - - # Split the text by existing newlines first - paragraphs = text.split("\n") - all_lines = [] - - for paragraph in paragraphs: - if not paragraph: - # Keep empty lines - all_lines.append("") - continue - - # Wrap each paragraph separately - wrapped = textwrap.wrap( - paragraph, - width=self.max_line_length, - break_long_words=True, - replace_whitespace=False, - drop_whitespace=True, + adjusted_properties[:, :4] = snap_boxes( + adjusted_properties[:, :4], resolution_wh ) - # Add the wrapped lines for this paragraph - if wrapped: - all_lines.extend(wrapped) - else: - # If wrap returns an empty list (e.g., for whitespace-only input) - all_lines.append("") - - return all_lines if all_lines else [""] + return adjusted_properties class BoxAnnotator(BaseAnnotator): @@ -1301,12 +1107,12 @@ class LabelAnnotator(_BaseLabelAnnotator): custom_color_lookup: Optional[np.ndarray] = None, ) -> np.ndarray: assert isinstance(scene, np.ndarray) - self._validate_labels(labels, detections) + validate_labels(labels, detections) - labels = self._get_labels_text(detections, labels) + labels = get_labels_text(detections, labels) label_properties = self._get_label_properties( - scene, detections, labels - ) # Pass scene + detections, labels + ) if self.smart_position: xyxy = label_properties[:, :4] @@ -1314,13 +1120,9 @@ class LabelAnnotator(_BaseLabelAnnotator): label_properties[:, :4] = xyxy label_properties = self._adjust_labels_in_frame( - scene.shape[1], - scene.shape[0], + (scene.shape[1], scene.shape[0]), labels, label_properties, - lambda text: cv2.getTextSize( - text, CV2_FONT, self.text_scale, self.text_thickness - )[0][0], ) self._draw_labels( @@ -1335,7 +1137,6 @@ class LabelAnnotator(_BaseLabelAnnotator): def _get_label_properties( self, - scene: np.ndarray, # Receive scene detections: Detections, labels: List[str], ) -> np.ndarray: @@ -1345,7 +1146,7 @@ class LabelAnnotator(_BaseLabelAnnotator): ).astype(int) for label, center_coords in zip(labels, anchors_coordinates): - wrapped_lines = self._wrap_text(label) + wrapped_lines = wrap_text(label, self.max_line_length) line_heights = [] line_widths = [] @@ -1375,12 +1176,10 @@ class LabelAnnotator(_BaseLabelAnnotator): position=self.text_anchor, ) - label_properties.append( - [ - *text_background_xyxy, - total_height, - ] - ) + label_properties.append([ + *text_background_xyxy, + total_height, + ]) return np.array(label_properties).reshape(-1, 5) def _draw_labels( @@ -1427,7 +1226,7 @@ class LabelAnnotator(_BaseLabelAnnotator): ) # Handle multiline text - wrapped_lines = self._wrap_text(labels[idx]) + wrapped_lines = wrap_text(labels[idx], self.max_line_length) current_y = box_xyxy[1] + self.text_padding # Start y position for line in wrapped_lines: @@ -1550,10 +1349,10 @@ class RichLabelAnnotator(_BaseLabelAnnotator): custom_color_lookup: Optional[np.ndarray] = None, ) -> ImageType: assert isinstance(scene, Image.Image) - self._validate_labels(labels, detections) + validate_labels(labels, detections) draw = ImageDraw.Draw(scene) - labels = self._get_labels_text(detections, labels) + labels = get_labels_text(detections, labels) label_properties = self._get_label_properties(draw, detections, labels) if self.smart_position: @@ -1562,14 +1361,9 @@ class RichLabelAnnotator(_BaseLabelAnnotator): label_properties[:, :4] = xyxy label_properties = self._adjust_labels_in_frame( - scene.width, - scene.height, + (scene.width, scene.height), labels, label_properties, - lambda text: int( - draw.textbbox((0, 0), text, font=self.font)[2] - - draw.textbbox((0, 0), text, font=self.font)[0] - ), ) self._draw_labels( @@ -1592,7 +1386,7 @@ class RichLabelAnnotator(_BaseLabelAnnotator): ).astype(int) for label, center_coords in zip(labels, anchor_coordinates): - wrapped_lines = self._wrap_text(label) + wrapped_lines = wrap_text(label, self.max_line_length) # Calculate the total text height and maximum width max_width = 0 @@ -1672,7 +1466,7 @@ class RichLabelAnnotator(_BaseLabelAnnotator): ) # Draw each line of text - wrapped_lines = self._wrap_text(labels[idx]) + wrapped_lines = wrap_text(labels[idx], self.max_line_length) x_position = box_xyxy[0] + self.text_padding - text_left y_position = box_xyxy[1] + self.text_padding - text_top diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index ae4f8cdc..48ee1ecb 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -1,8 +1,10 @@ +import textwrap from enum import Enum from typing import Optional, Tuple, Union import numpy as np +from supervision.config import CLASS_NAME_DATA_FIELD from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette from supervision.geometry.core import Position @@ -139,6 +141,178 @@ def resolve_color( return get_color_by_index(color=color, idx=idx) +def wrap_text(text: str, max_line_length=None) -> list[str]: + """ + Wraps text to the specified maximum line length, respecting existing newlines. + Uses the textwrap library for robust text wrapping. + + Args: + text (str): The text to wrap. + + Returns: + List[str]: A list of text lines after wrapping. + """ + + if not text: + return [""] + + if max_line_length is None: + return text.splitlines() or [""] + + paragraphs = text.split("\n") + all_lines = [] + + for paragraph in paragraphs: + if not paragraph: + # Keep empty lines + all_lines.append("") + continue + + wrapped = textwrap.wrap( + paragraph, + width=max_line_length, + break_long_words=True, + replace_whitespace=False, + drop_whitespace=True, + ) + + if wrapped: + all_lines.extend(wrapped) + else: + all_lines.append("") + + return all_lines if all_lines else [""] + + +def validate_labels(labels: Optional[list[str]], detections: Detections): + """ + Validates that the number of provided labels matches the number of detections. + + Args: + labels (Optional[List[str]]): A list of labels, one for each detection. Can + be None. + detections (Detections): The detections to be labeled. + + Raises: + ValueError: If `labels` is not None and its length does not match 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 " + f"number of detections ({len(detections)}). Each detection " + f"should have exactly 1 label." + ) + + +def get_labels_text( + detections: Detections, custom_labels: Optional[list[str]] +) -> list[str]: + """ + Retrieves the text labels for the detections. + + If `custom_labels` are provided, they are used. Otherwise, the labels are + extracted from the `detections` object, prioritizing the 'class_name' field, + then the `class_id`, and finally using the detection index as a string. + + Args: + detections (Detections): The detections to get labels for. + custom_labels (Optional[List[str]]): An optional list of custom labels. + + Returns: + List[str]: A list of text labels for each detection. + """ + 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 snap_boxes(xyxy: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray: + """ + Shifts bounding boxes into the frame so that they are fully contained + within the given resolution, prioritizing the top/left edge. + Unlike `clip_boxes`, this function does not crop boxes. + It moves them entirely if they exceed the frame boundaries. + + Args: + xyxy (np.ndarray): A numpy array of shape `(N, 4)` where each + row corresponds to a bounding box in the format + `(x_min, y_min, x_max, y_max)`. + resolution_wh (Tuple[int, int]): A tuple `(width, height)` + representing the resolution of the frame. + + Returns: + np.ndarray: A numpy array of shape `(N, 4)` with boxes shifted into frame. + + Examples: + ```python + import numpy as np + # Assuming this function is part of a library like supervision (sv) + # import supervision as sv # or just use the function directly + + xyxy = np.array([ + [-10, 10, 30, 50], # Off left + [310, 200, 350, 250], # Off right + [100, -20, 150, 30], # Off top + [200, 220, 250, 270], # Off bottom + [-20, 10, 350, 50], # Wider than frame, (width = 370 vs 320) + [10, -20, 30, 260] # Taller than frame, (height = 280 vs 240) + ]) + + resolution_wh = (320, 240) + + # Expected output for the new cases: + # [-20, 10, 350, 50] (wider) -> shifted right by -(-20) = 20 -> [0, 10, 370, 50] + # [10, -20, 30, 260] (taller) -> shifted down by -(-20) = 20 -> [10, 0, 30, 280] + # Note: Oversized boxes still won't be fully contained without cropping + # but this logic ensures the primary (top/left) boundary is corrected. + + snapped_boxes = snap_boxes(xyxy=xyxy, resolution_wh=resolution_wh) + print(snapped_boxes) + # Expected output (including original examples and new ones): + # [[ 0 10 40 50] # Original example 1 snapped + # [280 200 320 250] # Original example 2 snapped + # [100 0 150 50] # Original example 3 snapped + # [200 190 250 240] # Original example 4 snapped + # [ 0 10 370 50] # New example (wider) snapped by left edge priority + # [ 10 0 30 280]] # New example (taller) snapped by top edge priority + ``` + """ + result = np.copy(xyxy) + width, height = resolution_wh + + shift_if_left_out = -result[:, 0] + shift_if_right_out = width - result[:, 2] + + shift_x = np.where(result[:, 0] < 0, shift_if_left_out, + np.where(result[:, 2] > width, shift_if_right_out, 0)) + + result[:, 0] += shift_x + result[:, 2] += shift_x + + + shift_if_top_out = -result[:, 1] + shift_if_bottom_out = height - result[:, 3] + + shift_y = np.where(result[:, 1] < 0, shift_if_top_out, + np.where(result[:, 3] > height, shift_if_bottom_out, 0)) + + + result[:, 1] += shift_y + result[:, 3] += shift_y + + return result + + class Trace: def __init__( self, @@ -157,9 +331,10 @@ class Trace: def put(self, detections: Detections) -> None: frame_id = np.full(len(detections), self.current_frame_id, dtype=int) self.frame_id = np.concatenate([self.frame_id, frame_id]) - self.xy = np.concatenate( - [self.xy, detections.get_anchors_coordinates(self.anchor)] - ) + self.xy = np.concatenate([ + self.xy, + detections.get_anchors_coordinates(self.anchor), + ]) self.tracker_id = np.concatenate([self.tracker_id, detections.tracker_id]) unique_frame_id = np.unique(self.frame_id) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 3df4fad2..8949537e 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1264,12 +1264,10 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: Returns: Array of cross products of shape (number of anchors, detections) """ - vector_at_zero = np.array( - [ - vector.end.x - vector.start.x, - vector.end.y - vector.start.y, - ] - ) + vector_at_zero = np.array([ + vector.end.x - vector.start.x, + vector.end.y - vector.start.y, + ]) vector_start = np.array([vector.start.x, vector.start.y]) return np.cross(vector_at_zero, anchors - vector_start) @@ -1277,8 +1275,6 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: def spread_out_boxes( xyxy: np.ndarray, max_iterations: int = 100, - force_scale: float = 10.0, - consider_size: bool = True, ) -> np.ndarray: """ Spread out boxes that overlap with each other. @@ -1286,50 +1282,29 @@ 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_scale: Scale factor for the repulsion forces. - consider_size: Whether to consider box size when calculating forces. """ if len(xyxy) == 0: return xyxy xyxy_padded = pad_boxes(xyxy, px=1) - - # Calculate box areas if we're considering size - if consider_size: - box_areas = (xyxy_padded[:, 2] - xyxy_padded[:, 0]) * ( - xyxy_padded[:, 3] - xyxy_padded[:, 1] - ) - # Calculate the size factors (normalize by mean size) - size_factors = ( - np.sqrt(box_areas) / np.mean(np.sqrt(box_areas)) - if len(box_areas) > 0 - else np.ones(len(xyxy_padded)) - ) - size_factors = np.clip(size_factors, 0.5, 2.0) # Clip to avoid extreme values - else: - size_factors = np.ones(len(xyxy_padded)) - for _ in range(max_iterations): - # NxN matrix of IoU values + # NxN iou = box_iou_batch(xyxy_padded, xyxy_padded) - np.fill_diagonal(iou, 0) # Eliminate self-interactions - + np.fill_diagonal(iou, 0) if np.all(iou == 0): - break # No more overlaps + break overlap_mask = iou > 0 - # Nx2 array of box centers + # Nx2 centers = (xyxy_padded[:, :2] + xyxy_padded[:, 2:]) / 2 - # NxNx2 array of differences between centers + # NxNx2 delta_centers = centers[:, np.newaxis, :] - centers[np.newaxis, :, :] delta_centers *= overlap_mask[:, :, np.newaxis] - # Nx2 array of total delta vectors + # Nx2 delta_sum = np.sum(delta_centers, axis=1) - - # Normalize to unit vectors delta_magnitude = np.linalg.norm(delta_sum, axis=1, keepdims=True) direction_vectors = np.divide( delta_sum, @@ -1338,32 +1313,15 @@ def spread_out_boxes( where=delta_magnitude != 0, ) - # Calculate force based on IoU values force_vectors = np.sum(iou, axis=1) force_vectors = force_vectors[:, np.newaxis] * direction_vectors - # Apply size-based scaling if enabled - if consider_size: - force_vectors *= size_factors[:, np.newaxis] + force_vectors *= 10 + force_vectors[(force_vectors > 0) & (force_vectors < 2)] = 2 + force_vectors[(force_vectors < 0) & (force_vectors > -2)] = -2 - # Scale forces - force_vectors *= force_scale - - # Ensure minimum force for small overlaps - force_magnitudes = np.linalg.norm(force_vectors, axis=1, keepdims=True) - small_force_mask = (force_magnitudes > 0) & (force_magnitudes < 2) - if np.any(small_force_mask): - force_directions = force_vectors / np.where( - force_magnitudes > 0, force_magnitudes, 1 - ) - force_vectors = np.where( - small_force_mask, force_directions * 2, force_vectors - ) - - # Convert to integer displacements force_vectors = force_vectors.astype(int) - # Apply forces to update box positions xyxy_padded[:, [0, 1]] += force_vectors xyxy_padded[:, [2, 3]] += force_vectors From a6aa9b7d3be8099c158a95a7a9324412d86ef059 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 24 Apr 2025 23:33:11 +0000 Subject: [PATCH 12/15] =?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 | 14 +++++++------- supervision/annotators/utils.py | 26 ++++++++++++++++---------- supervision/detection/utils.py | 10 ++++++---- 3 files changed, 29 insertions(+), 21 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 5190fa21..f10491ff 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1110,9 +1110,7 @@ class LabelAnnotator(_BaseLabelAnnotator): validate_labels(labels, detections) labels = get_labels_text(detections, labels) - label_properties = self._get_label_properties( - detections, labels - ) + label_properties = self._get_label_properties(detections, labels) if self.smart_position: xyxy = label_properties[:, :4] @@ -1176,10 +1174,12 @@ class LabelAnnotator(_BaseLabelAnnotator): position=self.text_anchor, ) - label_properties.append([ - *text_background_xyxy, - total_height, - ]) + label_properties.append( + [ + *text_background_xyxy, + total_height, + ] + ) return np.array(label_properties).reshape(-1, 5) def _draw_labels( diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 48ee1ecb..8406bb5d 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -293,19 +293,23 @@ def snap_boxes(xyxy: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray: shift_if_left_out = -result[:, 0] shift_if_right_out = width - result[:, 2] - shift_x = np.where(result[:, 0] < 0, shift_if_left_out, - np.where(result[:, 2] > width, shift_if_right_out, 0)) + shift_x = np.where( + result[:, 0] < 0, + shift_if_left_out, + np.where(result[:, 2] > width, shift_if_right_out, 0), + ) result[:, 0] += shift_x result[:, 2] += shift_x - shift_if_top_out = -result[:, 1] shift_if_bottom_out = height - result[:, 3] - shift_y = np.where(result[:, 1] < 0, shift_if_top_out, - np.where(result[:, 3] > height, shift_if_bottom_out, 0)) - + shift_y = np.where( + result[:, 1] < 0, + shift_if_top_out, + np.where(result[:, 3] > height, shift_if_bottom_out, 0), + ) result[:, 1] += shift_y result[:, 3] += shift_y @@ -331,10 +335,12 @@ class Trace: def put(self, detections: Detections) -> None: frame_id = np.full(len(detections), self.current_frame_id, dtype=int) self.frame_id = np.concatenate([self.frame_id, frame_id]) - self.xy = np.concatenate([ - self.xy, - detections.get_anchors_coordinates(self.anchor), - ]) + self.xy = np.concatenate( + [ + self.xy, + detections.get_anchors_coordinates(self.anchor), + ] + ) self.tracker_id = np.concatenate([self.tracker_id, detections.tracker_id]) unique_frame_id = np.unique(self.frame_id) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 8949537e..8f1019e9 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1264,10 +1264,12 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: Returns: Array of cross products of shape (number of anchors, detections) """ - vector_at_zero = np.array([ - vector.end.x - vector.start.x, - vector.end.y - vector.start.y, - ]) + vector_at_zero = np.array( + [ + vector.end.x - vector.start.x, + vector.end.y - vector.start.y, + ] + ) vector_start = np.array([vector.start.x, vector.start.y]) return np.cross(vector_at_zero, anchors - vector_start) From 0c2aae6bd2721f473fe2ae34a3ba02c456aba1d6 Mon Sep 17 00:00:00 2001 From: hidara2000 Date: Fri, 25 Apr 2025 09:48:47 +1000 Subject: [PATCH 13/15] refactor(snap_boxes): enhance docstring and improve bounding box shifting logic --- supervision/annotators/utils.py | 84 +++++++++++++-------------------- 1 file changed, 34 insertions(+), 50 deletions(-) diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 8406bb5d..37967520 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -238,7 +238,7 @@ def get_labels_text( def snap_boxes(xyxy: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray: """ - Shifts bounding boxes into the frame so that they are fully contained + Shifts `label` bounding boxes into the frame so that they are fully contained within the given resolution, prioritizing the top/left edge. Unlike `clip_boxes`, this function does not crop boxes. It moves them entirely if they exceed the frame boundaries. @@ -254,65 +254,49 @@ def snap_boxes(xyxy: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray: np.ndarray: A numpy array of shape `(N, 4)` with boxes shifted into frame. Examples: - ```python - import numpy as np - # Assuming this function is part of a library like supervision (sv) - # import supervision as sv # or just use the function directly + ```python + import numpy as np - xyxy = np.array([ - [-10, 10, 30, 50], # Off left - [310, 200, 350, 250], # Off right - [100, -20, 150, 30], # Off top - [200, 220, 250, 270], # Off bottom - [-20, 10, 350, 50], # Wider than frame, (width = 370 vs 320) - [10, -20, 30, 260] # Taller than frame, (height = 280 vs 240) - ]) + # Example boxes: + xyxy = np.array([ + [-10, 10, 30, 50], # Off left edge + [310, 200, 350, 250], # Off right edge + [100, -20, 150, 30], # Off top edge + [200, 220, 250, 270], # Off bottom edge + [-20, 10, 350, 50], # Wider than frame (370 vs 320) + [10, -20, 30, 260] # Taller than frame (280 vs 240) + ]) - resolution_wh = (320, 240) + resolution_wh = (320, 240) + snapped_boxes = snap_boxes(xyxy=xyxy, resolution_wh=resolution_wh) - # Expected output for the new cases: - # [-20, 10, 350, 50] (wider) -> shifted right by -(-20) = 20 -> [0, 10, 370, 50] - # [10, -20, 30, 260] (taller) -> shifted down by -(-20) = 20 -> [10, 0, 30, 280] - # Note: Oversized boxes still won't be fully contained without cropping - # but this logic ensures the primary (top/left) boundary is corrected. - - snapped_boxes = snap_boxes(xyxy=xyxy, resolution_wh=resolution_wh) - print(snapped_boxes) - # Expected output (including original examples and new ones): - # [[ 0 10 40 50] # Original example 1 snapped - # [280 200 320 250] # Original example 2 snapped - # [100 0 150 50] # Original example 3 snapped - # [200 190 250 240] # Original example 4 snapped - # [ 0 10 370 50] # New example (wider) snapped by left edge priority - # [ 10 0 30 280]] # New example (taller) snapped by top edge priority - ``` + # Results: + # [[ 0 10 40 50] # Left edge shifted right + # [280 200 320 250] # Right edge shifted left + # [100 0 150 50] # Top edge shifted down + # [200 190 250 240] # Bottom edge shifted up + # [ 0 10 370 50] # Wide box aligned to left edge + # [ 10 0 30 280]] # Tall box aligned to top edge + ``` """ result = np.copy(xyxy) width, height = resolution_wh - shift_if_left_out = -result[:, 0] - shift_if_right_out = width - result[:, 2] + # X-axis (prioritize left edge) + left_overflow = result[:, 0] < 0 + result[left_overflow, 0:3:2] -= result[left_overflow, 0:1] - shift_x = np.where( - result[:, 0] < 0, - shift_if_left_out, - np.where(result[:, 2] > width, shift_if_right_out, 0), - ) + right_overflow = (~left_overflow) & (result[:, 2] > width) + right_shift = width - result[right_overflow, 2] + result[right_overflow, 0:3:2] += right_shift[:, np.newaxis] - result[:, 0] += shift_x - result[:, 2] += shift_x + # Y-axis (prioritize top edge) + top_overflow = result[:, 1] < 0 + result[top_overflow, 1:4:2] -= result[top_overflow, 1:2] - shift_if_top_out = -result[:, 1] - shift_if_bottom_out = height - result[:, 3] - - shift_y = np.where( - result[:, 1] < 0, - shift_if_top_out, - np.where(result[:, 3] > height, shift_if_bottom_out, 0), - ) - - result[:, 1] += shift_y - result[:, 3] += shift_y + bottom_overflow = (~top_overflow) & (result[:, 3] > height) + bottom_shift = height - result[bottom_overflow, 3] + result[bottom_overflow, 1:4:2] += bottom_shift[:, np.newaxis] return result From d7e0dedf74c0556693dc8c059a181111edb8ea89 Mon Sep 17 00:00:00 2001 From: hidara2000 Date: Fri, 11 Jul 2025 09:57:06 +1000 Subject: [PATCH 14/15] pull from main --- 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 f10491ff..b56208f4 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -139,10 +139,10 @@ class _BaseLabelAnnotator(BaseAnnotator): # Update the properties with the spread out boxes adjusted_properties[:, :4] = spread_boxes - # Additional check to ensure boxes are still within frame after spreading - adjusted_properties[:, :4] = snap_boxes( - adjusted_properties[:, :4], resolution_wh - ) + # Additional check to ensure boxes are still within frame after spreading + adjusted_properties[:, :4] = snap_boxes( + adjusted_properties[:, :4], resolution_wh + ) return adjusted_properties From 0b00edc526ba43369e36b98456e5dbb2d07dd5c6 Mon Sep 17 00:00:00 2001 From: hidara2000 Date: Fri, 11 Jul 2025 10:14:07 +1000 Subject: [PATCH 15/15] Remove redundant duplicated calls to spread_out_boxes --- 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 b56208f4..f10491ff 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -139,10 +139,10 @@ class _BaseLabelAnnotator(BaseAnnotator): # Update the properties with the spread out boxes adjusted_properties[:, :4] = spread_boxes - # Additional check to ensure boxes are still within frame after spreading - adjusted_properties[:, :4] = snap_boxes( - adjusted_properties[:, :4], resolution_wh - ) + # Additional check to ensure boxes are still within frame after spreading + adjusted_properties[:, :4] = snap_boxes( + adjusted_properties[:, :4], resolution_wh + ) return adjusted_properties