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

Add Pillow dependency, add docsrings for ImageType

Browse Source 
* PR #875
This commit is contained in:
Linas Kondrackis committed 2024-02-14 13:52:15 +02:00
1 parent c8f7622e91
commit 08418246c2
6 files changed
+139 -95
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

poetry.lock
Generated
+1 -13
View File
@@ -4105,18 +4105,6 @@ rfc3986 = ">=1.4.0"
rich = ">=12.0.0"
urllib3 = ">=1.26.0"
[[package]]
name = "types-pillow"
version = "10.2.0.20240206"
description = "Typing stubs for Pillow"
category = "main"
optional = false
python-versions = ">=3.8"
files = [
{file = "types-Pillow-10.2.0.20240206.tar.gz", hash = "sha256:f0de5107ff8362ffdbbd53ec896202ac905e6ab22ae784b46bcdad160ea143b9"},
{file = "types_Pillow-10.2.0.20240206-py3-none-any.whl", hash = "sha256:abc339ae28af5916146a7729261480d68ac902cd4ff57e0bdd402eee7962644d"},
]
[[package]]
name = "types-python-dateutil"
version = "2.8.19.20240106"
@@ -4356,4 +4344,4 @@ desktop = ["opencv-python"]
[metadata]
lock-version = "2.0"
python-versions = "^3.8"
content-hash = "fb7dad1d3b29d11473fcaad69324e04fa3b2ee6fade58830fc663c36008fffbd"
content-hash = "844f11f37895e0e56d43bd31f189e84e64c21248e120260a067597c822f35cc1"
pyproject.toml
+1 -1
View File
@@ -44,7 +44,7 @@ opencv-python = { version = ">=4.5.5.64", optional = true }
opencv-python-headless = ">=4.5.5.64"
requests = { version = ">=2.26.0,<=2.31.0", optional = true }
tqdm = { version = ">=4.62.3,<=4.66.1", optional = true }
types-pillow = "^10.2.0.20240206"
pillow = "^10.2.0"
[tool.poetry.extras]
desktop = ["opencv-python"]
supervision/annotators/base.py
+2 -2
View File
@@ -6,10 +6,10 @@ from PIL import Image
from supervision.detection.core import Detections
ImgType = TypeVar("ImgType", np.ndarray, Image.Image)
ImageType = TypeVar("ImageType", np.ndarray, Image.Image)
class BaseAnnotator(ABC):
@abstractmethod
def annotate(self, scene: ImgType, detections: Detections) -> ImgType:
def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
pass
supervision/annotators/core.py
+125 -72
View File
@@ -4,7 +4,7 @@ from typing import List, Optional, Tuple, Union
import cv2
import numpy as np
from supervision.annotators.base import BaseAnnotator, ImgType
from supervision.annotators.base import BaseAnnotator, ImageType
from supervision.annotators.utils import (
ColorLookup,
Trace,
@@ -45,21 +45,23 @@ class BoundingBoxAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with bounding boxes based on the provided detections.
Args:
scene (ImgType): The image where bounding boxes will be drawn.
scene (ImageType): The image where bounding boxes will be drawn. `ImageType`
is a flexible type, accepting either `numpy.ndarray` or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -124,21 +126,24 @@ class OrientedBoxAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with oriented bounding boxes based on the provided detections.
Args:
scene (ImgType): The image where bounding boxes will be drawn.
scene (ImageType): The image where bounding boxes will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -209,21 +214,24 @@ class MaskAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with masks based on the provided detections.
Args:
scene (ImgType): The image where masks will be drawn.
scene (ImageType): The image where masks will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -293,21 +301,24 @@ class PolygonAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with polygons based on the provided detections.
Args:
scene (ImgType): The image where polygons will be drawn.
scene (ImageType): The image where polygons will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -376,21 +387,24 @@ class ColorAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with box masks based on the provided detections.
Args:
scene (ImgType): The image where bounding boxes will be drawn.
scene (ImageType): The image where bounding boxes will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -467,21 +481,24 @@ class HaloAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with halos based on the provided detections.
Args:
scene (ImgType): The image where masks will be drawn.
scene (ImageType): The image where masks will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -562,21 +579,24 @@ class EllipseAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with ellipses based on the provided detections.
Args:
scene (ImgType): The image where ellipses will be drawn.
scene (ImageType): The image where ellipses will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -650,21 +670,24 @@ class BoxCornerAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with box corners based on the provided detections.
Args:
scene (ImgType): The image where box corners will be drawn.
scene (ImageType): The image where box corners will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -735,21 +758,24 @@ class CircleAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with circles based on the provided detections.
Args:
scene (ImgType): The image where box corners will be drawn.
scene (ImageType): The image where box corners will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -822,21 +848,24 @@ class DotAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with dots based on the provided detections.
Args:
scene (ImgType): The image where dots will be drawn.
scene (ImageType): The image where dots will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -948,23 +977,26 @@ class LabelAnnotator:
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
labels: List[str] = None,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with labels based on the provided detections.
Args:
scene (ImgType): The image where labels will be drawn.
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 (List[str]): Optional. 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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -1068,18 +1100,21 @@ class BlurAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene by blurring regions based on the provided detections.
Args:
scene (ImgType): The image where blurring will be applied.
scene (ImageType): The image where blurring will be applied.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
detections (Detections): Object detections to annotate.
Returns:
The annotated image.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -1150,22 +1185,25 @@ class TraceAnnotator:
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Draws trace paths on the frame based on the detection coordinates provided.
Args:
scene (ImgType): The image on which the traces will be drawn.
scene (ImageType): The image on which the traces will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
detections (Detections): The detections which include coordinates for
which the traces will be drawn.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
Returns:
The annotated image.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -1252,16 +1290,19 @@ class HeatMapAnnotator:
self.low_hue = low_hue
@scene_to_annotator_img_type
def annotate(self, scene: ImgType, detections: Detections) -> ImgType:
def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
"""
Annotates the scene with a heatmap based on the provided detections.
Args:
scene (ImgType): The image where the heatmap will be drawn.
scene (ImageType): The image where the heatmap will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
detections (Detections): Object detections to annotate.
Returns:
Annotated image.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -1327,19 +1368,22 @@ class PixelateAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene by pixelating regions based on the provided
detections.
Args:
scene (ImgType): The image where pixelating will be applied.
scene (ImageType): The image where pixelating will be applied.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
detections (Detections): Object detections to annotate.
Returns:
The annotated image.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -1412,21 +1456,24 @@ class TriangleAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with triangles based on the provided detections.
Args:
scene (ImgType): The image where triangles will be drawn.
scene (ImageType): The image where triangles will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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:
np.ndarray: The annotated image.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -1505,22 +1552,25 @@ class RoundBoxAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with bounding boxes with rounded edges
based on the provided detections.
Args:
scene (ImgType): The image where rounded bounding boxes will be drawn.
scene (ImageType): The image where rounded bounding boxes will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
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.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
@@ -1639,18 +1689,20 @@ class PercentageBarAnnotator(BaseAnnotator):
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
custom_color_lookup: Optional[np.ndarray] = None,
custom_values: Optional[np.ndarray] = None,
) -> ImgType:
) -> ImageType:
"""
Annotates the given scene with percentage bars based on the provided
detections. The percentage bars visually represent the confidence or custom
values associated with each detection.
Args:
scene (ImgType): The image where percentage bars will be drawn.
scene (ImageType): The image where percentage bars will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
detections (Detections): Object detections to annotate.
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
Allows to override the default color mapping strategy.
@@ -1661,7 +1713,8 @@ class PercentageBarAnnotator(BaseAnnotator):
to be displayed.
Returns:
The annotated image.
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)
Example:
```python
supervision/annotators/utils.py
+2 -2
View File
@@ -6,7 +6,7 @@ import cv2
import numpy as np
from PIL import Image
from supervision.annotators.base import ImgType
from supervision.annotators.base import ImageType
from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette
from supervision.geometry.core import Position
@@ -133,7 +133,7 @@ def scene_to_annotator_img_type(annotate_func):
"""
@wraps(annotate_func)
def wrapper(self, scene: ImgType, *args, **kwargs):
def wrapper(self, scene: ImageType, *args, **kwargs):
if isinstance(scene, np.ndarray):
return annotate_func(self, scene, *args, **kwargs)
supervision/detection/annotate.py
+8 -5
View File
@@ -2,7 +2,7 @@ from typing import List, Optional, Union
import cv2
from supervision.annotators.base import ImgType
from supervision.annotators.base import ImageType
from supervision.annotators.utils import scene_to_annotator_img_type
from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette
@@ -49,16 +49,18 @@ class BoxAnnotator:
@scene_to_annotator_img_type
def annotate(
self,
scene: ImgType,
scene: ImageType,
detections: Detections,
labels: Optional[List[str]] = None,
skip_label: bool = False,
) -> ImgType:
) -> ImageType:
"""
Draws bounding boxes on the frame using the detections provided.
Args:
scene (ImgType): The image on which the bounding boxes will be drawn
scene (ImageType): The image on which the bounding boxes will be drawn.
`ImageType` is a flexible type, accepting either `numpy.ndarray`
or `PIL.Image.Image`.
detections (Detections): The detections for which the
bounding boxes will be drawn
labels (Optional[List[str]]): An optional list of labels
@@ -66,7 +68,8 @@ class BoxAnnotator:
corresponding `class_id` will be used as label.
skip_label (bool): Is set to `True`, skips bounding box label annotation.
Returns:
ImgType: The image with the bounding boxes drawn on it
ImageType: The image with the bounding boxes drawn on it, matching the
type of `scene` (`numpy.ndarray` or `PIL.Image.Image`)
Example:
```python