diff --git a/poetry.lock b/poetry.lock index c04982f0..2e001a93 100644 --- a/poetry.lock +++ b/poetry.lock @@ -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" diff --git a/pyproject.toml b/pyproject.toml index 1a6cc2ea..29207737 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/supervision/annotators/base.py b/supervision/annotators/base.py index 1eb66d33..baea60f3 100644 --- a/supervision/annotators/base.py +++ b/supervision/annotators/base.py @@ -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 diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 00e4c864..1186402b 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -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 diff --git a/supervision/annotators/utils.py b/supervision/annotators/utils.py index 45282760..ed0070f0 100644 --- a/supervision/annotators/utils.py +++ b/supervision/annotators/utils.py @@ -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) diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index b64df8e8..a8df0a57 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -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