📘 docs updates for `LabelAnnotator`
This commit is contained in:
SkalskiP
parent
560a35b75f
commit
3e8906d830
|
|
@ -13,3 +13,7 @@
|
|||
## BoxCornerAnnotator
|
||||
|
||||
:::supervision.annotators.core.BoxCornerAnnotator
|
||||
|
||||
## LabelAnnotator
|
||||
|
||||
:::supervision.annotators.core.LabelAnnotator
|
||||
|
|
@ -10,10 +10,8 @@ from supervision.annotators.core import (
|
|||
BoundingBoxAnnotator,
|
||||
BoxCornerAnnotator,
|
||||
EllipseAnnotator,
|
||||
LabelAdvancedAnnotator,
|
||||
LabelAnnotator,
|
||||
MaskAnnotator,
|
||||
TraceAnnotator,
|
||||
)
|
||||
from supervision.classification.core import Classifications
|
||||
from supervision.dataset.core import (
|
||||
|
|
|
|||
|
|
@ -307,6 +307,9 @@ class BoxCornerAnnotator(BaseAnnotator):
|
|||
|
||||
|
||||
class LabelAnnotator:
|
||||
"""
|
||||
A class for annotating labels on an image using provided detections.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
|
|
@ -317,6 +320,19 @@ class LabelAnnotator:
|
|||
text_position: Position = Position.TOP_LEFT,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating the text background.
|
||||
text_color (Color): The color 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_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.text_color: Color = text_color
|
||||
self.text_scale: float = text_scale
|
||||
|
|
@ -367,6 +383,31 @@ class LabelAnnotator:
|
|||
detections: Detections,
|
||||
labels: List[str] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with labels based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where labels will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
labels (List[str]): Optional. Custom labels for each detection.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> label_annotator = sv.LabelAnnotator()
|
||||
>>> annotated_frame = label_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
"""
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
for detection_idx in range(len(detections)):
|
||||
detection_xyxy = detections.xyxy[detection_idx].astype(int)
|
||||
|
|
@ -416,213 +457,3 @@ class LabelAnnotator:
|
|||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class LabelAdvancedAnnotator(BaseAnnotator):
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
text_color: Color = Color.black(),
|
||||
text_padding: int = 20,
|
||||
color_by_track: bool = False,
|
||||
font: Optional[str] = None,
|
||||
font_size: Optional[int] = 15,
|
||||
):
|
||||
if font and os.path.exists(font):
|
||||
self.font = ImageFont.truetype(font, font_size)
|
||||
else:
|
||||
self.font = ImageFont.load_default()
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.text_color: Color = text_color
|
||||
self.text_padding: int = text_padding
|
||||
self.color_by_track = color_by_track
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
labels: Optional[List[str]] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Draws text on the frame using the detections provided and label.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image on which the bounding boxes will be drawn
|
||||
detections (Detections): The detections for which
|
||||
the bounding boxes will be drawn
|
||||
labels (Optional[List[str]]): An optional list of labels corresponding
|
||||
to each detection. If `labels` are not provided,
|
||||
corresponding `class_id` will be used as label.
|
||||
Returns:
|
||||
np.ndarray: The image with the bounding boxes drawn on it
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> classes = ['person', ...]
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> pil_label_annotator = sv.LabelAdvancedAnnotator()
|
||||
>>> labels = [
|
||||
... f"{classes[class_id]} {confidence:0.2f}"
|
||||
... for _, _, confidence, class_id, _
|
||||
... in detections
|
||||
... ]
|
||||
>>> annotated_frame = pil_label_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections,
|
||||
... labels=labels,
|
||||
... )
|
||||
```
|
||||
"""
|
||||
pil_image = Image.fromarray(scene)
|
||||
draw = ImageDraw.Draw(pil_image)
|
||||
text_color = "#fff"
|
||||
|
||||
for i in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
|
||||
if self.color_by_track:
|
||||
tracker_id = (
|
||||
detections.tracker_id[i]
|
||||
if detections.tracker_id is not None
|
||||
else None
|
||||
)
|
||||
idx = tracker_id if tracker_id is not None else i
|
||||
else:
|
||||
class_id = (
|
||||
detections.class_id[i] if detections.class_id is not None else None
|
||||
)
|
||||
idx = class_id if class_id is not None else i
|
||||
|
||||
color = (
|
||||
self.color.by_idx(idx)
|
||||
if isinstance(self.color, ColorPalette)
|
||||
else self.color
|
||||
)
|
||||
|
||||
text = (
|
||||
f"{idx}"
|
||||
if (labels is None or len(detections) != len(labels))
|
||||
else labels[i]
|
||||
)
|
||||
|
||||
text_bbox = draw.textbbox((x1, y1), text, font=self.font)
|
||||
|
||||
text_height = text_bbox[3] - text_bbox[1]
|
||||
text_width = text_bbox[2] - text_bbox[0]
|
||||
|
||||
text_x = x1 + self.text_padding / 2
|
||||
text_y = y1 - self.text_padding / 2 - text_height
|
||||
|
||||
text_background_x1 = x1
|
||||
text_background_y1 = y1 - self.text_padding / 2 - text_height
|
||||
|
||||
text_background_x2 = x1 + 2 * self.text_padding / 2 + text_width
|
||||
text_background_y2 = y1 # correct
|
||||
|
||||
draw.rectangle(
|
||||
(
|
||||
text_background_x1,
|
||||
text_background_y1,
|
||||
text_background_x2,
|
||||
text_background_y2,
|
||||
),
|
||||
fill=color.as_bgr(),
|
||||
)
|
||||
draw.text((text_x, text_y), text, font=self.font, fill=text_color)
|
||||
|
||||
scene = np.asarray(pil_image)
|
||||
return scene
|
||||
|
||||
|
||||
class TraceAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing trajectory of a tracker on an image using detections provided.
|
||||
|
||||
Attributes:
|
||||
color (Union[Color, ColorPalette]): The color to draw the trajectory,
|
||||
can be a single color or a color palette
|
||||
color_by_track (bool): Whther to use tracker id to pick the color
|
||||
position (Optional[Position]): Choose position of trajectory such as
|
||||
center position, top left corner, etc
|
||||
trace_length (int): Length of the previous points
|
||||
thickness (int): thickness of the line
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
color_by_track: bool = False,
|
||||
position: Optional[Position] = Position.CENTER,
|
||||
trace_length: int = 30,
|
||||
thickness: int = 2,
|
||||
):
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.color_by_track = color_by_track
|
||||
self.position = position
|
||||
self.tracker_storage = defaultdict(lambda: [])
|
||||
self.trace_length = trace_length
|
||||
self.thickness = thickness
|
||||
|
||||
def annotate(
|
||||
self, scene: np.ndarray, detections: Detections, **kwargs
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Draw the object trajectory based on history of tracked objects
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image on which the trace will be drawn
|
||||
detections (Detections): The detections for trajectory and points
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the masks overlaid
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> classes = ['person', ...]
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> trace_annotator = sv.TraceAnnotator()
|
||||
>>> annotated_frame = trace_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
"""
|
||||
if detections.tracker_id is None:
|
||||
return scene
|
||||
|
||||
anchor_points = detections.get_anchor_coordinates(anchor=self.position)
|
||||
|
||||
for i, tracker_id in enumerate(detections.tracker_id):
|
||||
track = self.tracker_storage[tracker_id]
|
||||
track.append((anchor_points[i][0], anchor_points[i][1]))
|
||||
if len(track) > self.trace_length:
|
||||
track.pop(0)
|
||||
points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
|
||||
|
||||
if self.color_by_track:
|
||||
idx = tracker_id if tracker_id is not None else i
|
||||
else:
|
||||
class_id = (
|
||||
detections.class_id[i] if detections.class_id is not None else None
|
||||
)
|
||||
idx = class_id if class_id is not None else i
|
||||
color = (
|
||||
self.color.by_idx(idx)
|
||||
if isinstance(self.color, ColorPalette)
|
||||
else self.color
|
||||
)
|
||||
cv2.polylines(
|
||||
scene,
|
||||
[points],
|
||||
isClosed=False,
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
|
||||
return scene
|
||||
|
|
|
|||
Loading…
Reference in New Issue