diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index b5395737..5060cf97 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -12,10 +12,14 @@ from supervision.annotators.utils import ( PENDING_TRACK_ID, 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, @@ -41,6 +45,109 @@ 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. + max_line_length (Optional[int]): Maximum number of characters per line before + wrapping the text. None means no wrapping. + """ + + 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, + max_line_length: Optional[int] = None, + ): + """ + Initializes the _BaseLabelAnnotator. + + Args: + color (Union[Color, ColorPalette], optional): The color to use for the label + background. + text_color (Union[Color, ColorPalette], optional): The color to use for the + label text. + text_padding (int, optional): The padding around the label text, in pixels. + text_position (Position, optional): The position of the text relative to the + detection bounding box. + color_lookup (ColorLookup, optional): The method used to determine the color + of the label + border_radius (int, optional): The radius of the label background corners, + in pixels. + smart_position (bool, optional): Whether to intelligently adjust the label + position to avoid overlapping with other elements. + max_line_length (Optional[int], optional): Maximum number of characters per + line before wrapping the text. None means no wrapping. + + """ + 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.max_line_length: Optional[int] = max_line_length + + def _adjust_labels_in_frame( + self, + resolution_wh: Tuple[int, int], + labels: List[str], + label_properties: np.ndarray, + ) -> np.ndarray: + """ + Adjusts the position of labels to ensure they stay within the frame boundaries. + + Args: + frame_width (int): The width 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, ...]. + + Returns: + np.ndarray: The adjusted label properties. + """ + adjusted_properties = label_properties.copy() + + # First, make sure the boxes don't go outside the frame + adjusted_properties[:, :4] = snap_boxes( + adjusted_properties[:, :4], + resolution_wh, + ) + + # 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 + + # Additional check to ensure boxes are still within frame after spreading + adjusted_properties[:, :4] = snap_boxes( + adjusted_properties[:, :4], resolution_wh + ) + + return adjusted_properties + + class BoxAnnotator(BaseAnnotator): """ A class for drawing bounding boxes on an image using provided detections. @@ -961,7 +1068,7 @@ class DotAnnotator(BaseAnnotator): return scene -class LabelAnnotator(BaseAnnotator): +class LabelAnnotator(_BaseLabelAnnotator): """ A class for annotating labels on an image using provided detections. """ @@ -977,87 +1084,33 @@ class LabelAnnotator(BaseAnnotator): color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, smart_position: bool = False, + max_line_length: Optional[int] = None, ): - """ - 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, + max_line_length=max_line_length, + ) @ensure_cv2_image_for_annotation def annotate( self, - scene: ImageType, + scene: ImageType, # 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) + validate_labels(labels, detections) - labels = self._get_labels_text(detections, labels) + labels = get_labels_text(detections, labels) label_properties = self._get_label_properties(detections, labels) if self.smart_position: @@ -1065,6 +1118,12 @@ class LabelAnnotator(BaseAnnotator): xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy + label_properties = self._adjust_labels_in_frame( + (scene.shape[1], scene.shape[0]), + labels, + label_properties, + ) + self._draw_labels( scene=scene, labels=labels, @@ -1075,41 +1134,40 @@ 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, 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 ).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 = wrap_text(label, self.max_line_length) + 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), @@ -1120,29 +1178,11 @@ class LabelAnnotator(BaseAnnotator): label_properties.append( [ *text_background_xyxy, - text_h, + total_height, ] ) - 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 [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 _draw_labels( self, scene: np.ndarray, @@ -1177,8 +1217,8 @@ class LabelAnnotator(BaseAnnotator): color_lookup=color_lookup, ) - box_xyxy = label_property[:4] - text_height_padded = label_property[4] + box_xyxy = label_property[:4].astype(int) + self.draw_rounded_rectangle( scene=scene, xyxy=box_xyxy, @@ -1186,18 +1226,44 @@ class LabelAnnotator(BaseAnnotator): border_radius=self.border_radius, ) - text_x = box_xyxy[0] + self.text_padding - text_y = box_xyxy[1] + self.text_padding + 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 = wrap_text(labels[idx], self.max_line_length) + current_y = box_xyxy[1] + self.text_padding # Start y position + + for line in wrapped_lines: + if not line: + # Use a character with ascenders and descenders as height reference + (_, text_h) = cv2.getTextSize( + text="Tg", + 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( @@ -1242,7 +1308,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. @@ -1259,32 +1325,21 @@ class RichLabelAnnotator(BaseAnnotator): color_lookup: ColorLookup = ColorLookup.CLASS, border_radius: int = 0, smart_position: bool = False, + max_line_length: Optional[int] = None, ): - """ - 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, + max_line_length=max_line_length, + ) @ensure_pil_image_for_annotation def annotate( @@ -1294,50 +1349,11 @@ class RichLabelAnnotator(BaseAnnotator): 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 - ) - ``` - - """ 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: @@ -1345,6 +1361,12 @@ class RichLabelAnnotator(BaseAnnotator): xyxy = spread_out_boxes(xyxy) label_properties[:, :4] = xyxy + label_properties = self._adjust_labels_in_frame( + (scene.width, scene.height), + labels, + label_properties, + ) + self._draw_labels( draw=draw, labels=labels, @@ -1355,26 +1377,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( @@ -1382,13 +1387,26 @@ class RichLabelAnnotator(BaseAnnotator): ).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 = text_width + 2 * self.text_padding - height_padded = text_height + 2 * self.text_padding + wrapped_lines = wrap_text(label, self.max_line_length) + + # 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), @@ -1396,30 +1414,16 @@ class RichLabelAnnotator(BaseAnnotator): 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) - @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, @@ -1450,24 +1454,35 @@ 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 - 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 = 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 + + 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]): diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 74d25a65..88953799 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 @@ -144,6 +146,166 @@ 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 `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. + + 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 + + # 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) + snapped_boxes = snap_boxes(xyxy=xyxy, resolution_wh=resolution_wh) + + # 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 + + # X-axis (prioritize left edge) + left_overflow = result[:, 0] < 0 + result[left_overflow, 0:3:2] -= result[left_overflow, 0:1] + + right_overflow = (~left_overflow) & (result[:, 2] > width) + right_shift = width - result[right_overflow, 2] + result[right_overflow, 0:3:2] += right_shift[:, np.newaxis] + + # Y-axis (prioritize top edge) + top_overflow = result[:, 1] < 0 + result[top_overflow, 1:4:2] -= result[top_overflow, 1:2] + + 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 + + class Trace: def __init__( self, @@ -163,7 +325,10 @@ class Trace: 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, + detections.get_anchors_coordinates(self.anchor), + ] ) self.tracker_id = np.concatenate([self.tracker_id, detections.tracker_id]) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 5a8e8844..f7999cdd 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1503,7 +1503,10 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: 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.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)