Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add PolygonAnnotator and update documentation #602

Merged
merged 4 commits into from
Nov 15, 2023
Merged

Add PolygonAnnotator and update documentation #602

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e40532c
Add PolygonAnnotator and update documentation
SkalskiP Nov 15, 2023
777e8d3
fix(pre_commit): 🎨 auto format pre-commit hooks
pre-commit-ci[bot] Nov 15, 2023
c1beb0f
Add PolygonAnnotator to supervision annotators
SkalskiP Nov 15, 2023
2ed4e16
Added a visual and code example for `PolygonAnnotator` in `supervisio…
SkalskiP Nov 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions docs/annotators.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,27 @@

</div>

=== "Polygon"

```python
>>> import supervision as sv

>>> image = ...
>>> detections = sv.Detections(...)

>>> polygon_annotator = sv.PolygonAnnotator()
>>> annotated_frame = polygon_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```

<div class="result" markdown>

![polygon-annotator-example](https://media.roboflow.com/supervision-annotator-examples/polygon-annotator-example-purple.png){ align=center width="800" }

</div>

=== "Label"

```python
Expand Down Expand Up @@ -304,6 +325,10 @@

:::supervision.annotators.core.MaskAnnotator

## PolygonAnnotator

:::supervision.annotators.core.PolygonAnnotator

## LabelAnnotator

:::supervision.annotators.core.LabelAnnotator
Expand Down
1 change: 1 addition & 0 deletions supervision/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
HeatMapAnnotator,
LabelAnnotator,
MaskAnnotator,
PolygonAnnotator,
TraceAnnotator,
)
from supervision.annotators.utils import ColorLookup
Expand Down
123 changes: 109 additions & 14 deletions supervision/annotators/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,9 @@
from supervision.annotators.base import BaseAnnotator
from supervision.annotators.utils import ColorLookup, Trace, resolve_color
from supervision.detection.core import Detections
from supervision.detection.utils import clip_boxes
from supervision.detection.utils import clip_boxes, mask_to_polygons
from supervision.draw.color import Color, ColorPalette
from supervision.draw.utils import draw_polygon
from supervision.geometry.core import Position


Expand Down Expand Up @@ -51,7 +52,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -93,6 +94,10 @@ def annotate(
class MaskAnnotator(BaseAnnotator):
"""
A class for drawing masks on an image using provided detections.

!!! warning

This annotator utilizes the `sv.Detections.mask`.
"""

def __init__(
Expand Down Expand Up @@ -129,7 +134,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -170,6 +175,92 @@ def annotate(
return scene


class PolygonAnnotator(BaseAnnotator):
"""
A class for drawing polygons on an image using provided detections.

!!! warning

This annotator utilizes the `sv.Detections.mask`.
"""

def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 2,
color_lookup: ColorLookup = ColorLookup.CLASS,
):
"""
Args:
color (Union[Color, ColorPalette]): The color or color palette to use for
annotating detections.
thickness (int): Thickness of the polygon lines.
color_lookup (str): Strategy for mapping colors to annotations.
Options are `INDEX`, `CLASS`, `TRACE`.
"""
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.color_lookup: ColorLookup = color_lookup

def annotate(
self,
scene: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Annotates the given scene with polygons based on the provided detections.

Args:
scene (np.ndarray): The image where polygons will be drawn.
detections (Detections): Object detections to annotate.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.

Returns:
The annotated image.

Example:
```python
>>> import supervision as sv

>>> image = ...
>>> detections = sv.Detections(...)

>>> polygon_annotator = sv.PolygonAnnotator()
>>> annotated_frame = polygon_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```

![polygon-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/polygon-annotator-example-purple.png)
"""
if detections.mask is None:
return scene

for detection_idx in range(len(detections)):
mask = detections.mask[detection_idx]
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=self.color_lookup
if custom_color_lookup is None
else custom_color_lookup,
)
for polygon in mask_to_polygons(mask=mask):
scene = draw_polygon(
scene=scene,
polygon=polygon,
color=color,
thickness=self.thickness,
)

return scene


class BoxMaskAnnotator(BaseAnnotator):
"""
A class for drawing box masks on an image using provided detections.
Expand Down Expand Up @@ -209,7 +300,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -255,6 +346,10 @@ def annotate(
class HaloAnnotator(BaseAnnotator):
"""
A class for drawing Halos on an image using provided detections.

!!! warning

This annotator utilizes the `sv.Detections.mask`.
"""

def __init__(
Expand Down Expand Up @@ -295,7 +390,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -389,7 +484,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -476,7 +571,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -560,7 +655,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -646,7 +741,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -776,7 +871,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -873,7 +968,7 @@ def annotate(
detections (Detections): Object detections to annotate.

Returns:
np.ndarray: The annotated image.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -911,7 +1006,7 @@ class TraceAnnotator:

!!! warning

This annotator utilizes the `tracker_id`. Read
This annotator utilizes the `sv.Detections.tracker_id`. Read
[here](https://supervision.roboflow.com/trackers/) to learn how to plug
tracking into your inference pipeline.
"""
Expand Down Expand Up @@ -959,7 +1054,7 @@ def annotate(
Allows to override the default color mapping strategy.

Returns:
np.ndarray: The image with the trace paths drawn on it.
The annotated image.

Example:
```python
Expand Down Expand Up @@ -1055,7 +1150,7 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
detections (Detections): Object detections to annotate.

Returns:
np.ndarray: Annotated image.
Annotated image.

Example:
```python
Expand Down