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

Merge pull request #422 from hardikdava/box_mask_annotator 

Added box mask annotators
This commit is contained in:
Piotr Skalski 2023-10-07 00:47:09 +02:00 committed by GitHub
parent 4733ebe27d 669c0658fe
commit d30d80717c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 153 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

133
docs/annotators.md
View File

@ -19,48 +19,6 @@
</div>
=== "Mask"
```python
>>> import supervision as sv
>>> image = ...
>>> detections = sv.Detections(...)
>>> mask_annotator = sv.MaskAnnotator()
>>> annotated_frame = mask_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
<div class="result" markdown>
![mask-annotator-example](https://media.roboflow.com/supervision-annotator-examples/mask-annotator-example.png){ align=center width="800" }
</div>
=== "Ellipse"
```python
>>> import supervision as sv
>>> image = ...
>>> detections = sv.Detections(...)
>>> ellipse_annotator = sv.EllipseAnnotator()
>>> annotated_frame = ellipse_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
<div class="result" markdown>
![ellipse-annotator-example](https://media.roboflow.com/supervision-annotator-examples/ellipse-annotator-example.png){ align=center width="800" }
</div>
=== "BoxCorner"
```python
@ -82,6 +40,27 @@
</div>
=== "BoxMaskAnnotator"
```python
>>> import supervision as sv
>>> image = ...
>>> detections = sv.Detections(...)
>>> box_mask_annotator = sv.BoxMaskAnnotator()
>>> annotated_frame = box_mask_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
<div class="result" markdown>
![box-mask-annotator-example](https://media.roboflow.com/supervision-annotator-examples/box-mask-annotator-example.png){ align=center width="800" }
</div>
=== "Circle"
```python
@ -103,6 +82,48 @@
</div>
=== "Ellipse"
```python
>>> import supervision as sv
>>> image = ...
>>> detections = sv.Detections(...)
>>> ellipse_annotator = sv.EllipseAnnotator()
>>> annotated_frame = ellipse_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
<div class="result" markdown>
![ellipse-annotator-example](https://media.roboflow.com/supervision-annotator-examples/ellipse-annotator-example.png){ align=center width="800" }
</div>
=== "Mask"
```python
>>> import supervision as sv
>>> image = ...
>>> detections = sv.Detections(...)
>>> mask_annotator = sv.MaskAnnotator()
>>> annotated_frame = mask_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
<div class="result" markdown>
![mask-annotator-example](https://media.roboflow.com/supervision-annotator-examples/mask-annotator-example.png){ align=center width="800" }
</div>
=== "Label"
```python
@ -170,30 +191,34 @@
:::supervision.annotators.core.BoundingBoxAnnotator
## MaskAnnotator
:::supervision.annotators.core.MaskAnnotator
## EllipseAnnotator
:::supervision.annotators.core.EllipseAnnotator
## BoxCornerAnnotator
:::supervision.annotators.core.BoxCornerAnnotator
## BoxMaskAnnotator
:::supervision.annotators.core.BoxMaskAnnotator
## CircleAnnotator
:::supervision.annotators.core.CircleAnnotator
## EllipseAnnotator
:::supervision.annotators.core.EllipseAnnotator
## MaskAnnotator
:::supervision.annotators.core.MaskAnnotator
## LabelAnnotator
:::supervision.annotators.core.LabelAnnotator
## TraceAnnotator
:::supervision.annotators.core.TraceAnnotator
## BlurAnnotator
:::supervision.annotators.core.BlurAnnotator
## TraceAnnotator
:::supervision.annotators.core.TraceAnnotator

1
supervision/__init__.py
View File

@ -10,6 +10,7 @@ from supervision.annotators.core import (
BlurAnnotator,
BoundingBoxAnnotator,
BoxCornerAnnotator,
BoxMaskAnnotator,
CircleAnnotator,
EllipseAnnotator,
LabelAnnotator,

73
supervision/annotators/core.py
View File

@ -156,6 +156,79 @@ class MaskAnnotator(BaseAnnotator):
return scene
class BoxMaskAnnotator(BaseAnnotator):
"""
A class for drawing box masks on an image using provided detections.
"""
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
opacity: float = 0.5,
color_map: str = "class",
):
"""
Args:
color (Union[Color, ColorPalette]): The color or color palette to use for
annotating detections.
opacity (float): Opacity of the overlay mask. Must be between `0` and `1`.
color_map (str): Strategy for mapping colors to annotations.
Options are `index`, `class`, or `track`.
"""
self.color: Union[Color, ColorPalette] = color
self.color_map: ColorMap = ColorMap(color_map)
self.opacity = opacity
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
"""
Annotates the given scene with box masks based on the provided detections.
Args:
scene (np.ndarray): The image where bounding boxes will be drawn.
detections (Detections): Object detections to annotate.
Returns:
np.ndarray: The annotated image.
Example:
```python
>>> import supervision as sv
>>> image = ...
>>> detections = sv.Detections(...)
>>> box_mask_annotator = sv.BoxMaskAnnotator()
>>> annotated_frame = box_mask_annotator.annotate(
... scene=image.copy(),
... detections=detections
... )
```
![box-mask-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/box-mask-annotator-example.png)
"""
mask_image = scene.copy()
for detection_idx in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
idx = resolve_color_idx(
detections=detections,
detection_idx=detection_idx,
color_map=self.color_map,
)
color = resolve_color(color=self.color, idx=idx)
cv2.rectangle(
img=scene,
pt1=(x1, y1),
pt2=(x2, y2),
color=color.as_bgr(),
thickness=-1,
)
scene = cv2.addWeighted(
scene, self.opacity, mask_image, 1 - self.opacity, gamma=0
)
return scene
class EllipseAnnotator(BaseAnnotator):
"""
A class for drawing ellipses on an image using provided detections.