okooooooo/supervision
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(annotators): streamline label handling and improve box snapping logic

Browse Source
This commit is contained in:
hidara2000 committed 2025-04-25 09:32:54 +10:00
1 parent 1876019faf
commit 29695055d2
3 files changed
+228 -301
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

supervision/annotators/core.py
+37 -243
View File
@@ -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
supervision/annotators/utils.py
+178 -3
View File
@@ -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)
supervision/detection/utils.py
+13 -55
View File
@@ -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