From 20b1b217741908c68ddea85d93e5041746948d10 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 28 Sep 2023 15:50:55 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=B9=20Removed=20old=20`MaskAnnotator`?= =?UTF-8?q?=20implementation=20as=20it=20is=20no=20longer=20used.=20Moved?= =?UTF-8?q?=20`ColorMap`,=20`resolve=5Fcolor=5Fidx`,=20`resolve=5Fcolor`?= =?UTF-8?q?=20and=20`BaseAnnotator`=20to=20`supervision/annotators/base`.?= =?UTF-8?q?=20Removed=20`BoxMaskAnnotator`=20as=20it=20can=20be=20achieved?= =?UTF-8?q?=20by=20`BoundingBoxAnnotator`=20and=20`MaskAnnotator`=20compos?= =?UTF-8?q?ition.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/__init__.py | 1 - supervision/annotators/base.py | 58 ++++++++++++++ supervision/annotators/core.py | 129 +----------------------------- supervision/detection/annotate.py | 58 -------------- 4 files changed, 60 insertions(+), 186 deletions(-) create mode 100644 supervision/annotators/base.py diff --git a/supervision/__init__.py b/supervision/__init__.py index a534484c..991f79a2 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -11,7 +11,6 @@ from supervision.annotators.composable import DetectionAnnotator, SegmentationAn from supervision.annotators.core import ( BoundingBoxAnnotator, BoxCornerAnnotator, - BoxMaskAnnotator, EllipseAnnotator, LabelAdvancedAnnotator, LabelAnnotator, diff --git a/supervision/annotators/base.py b/supervision/annotators/base.py new file mode 100644 index 00000000..d9a6e29a --- /dev/null +++ b/supervision/annotators/base.py @@ -0,0 +1,58 @@ +from abc import ABC, abstractmethod +from enum import Enum +from typing import Union + +import numpy as np + +from supervision.detection.core import Detections +from supervision.draw.color import Color, ColorPalette + + +class ColorMap(Enum): + """ + Enum for annotator color mapping. + """ + + INDEX = "index" + CLASS = "class" + TRACK = "track" + + +def resolve_color_idx( + detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS +) -> int: + if detection_idx >= len(detections): + raise ValueError( + f"Detection index {detection_idx}" + f"is out of bounds for detections of length {len(detections)}" + ) + + if color_map == ColorMap.INDEX: + return detection_idx + elif color_map == ColorMap.CLASS: + if detections.class_id is None: + raise ValueError( + "Could not resolve color by class because" + "Detections do not have class_id" + ) + return detections.class_id[detection_idx] + elif color_map == ColorMap.TRACK: + if detections.tracker_id is None: + raise ValueError( + "Could not resolve color by track because" + "Detections do not have tracker_id" + ) + return detections.tracker_id[detection_idx] + + +def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color: + if isinstance(color, ColorPalette): + return color.by_idx(idx) + else: + return color + + +class BaseAnnotator(ABC): + @abstractmethod + def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: + pass diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 9a4758b3..d7f4aff0 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,68 +1,18 @@ import os.path -from abc import ABC, abstractmethod from collections import defaultdict -from enum import Enum from typing import Callable, List, Optional, Union import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont +from supervision.annotators.base import BaseAnnotator, ColorMap, resolve_color_idx, \ + resolve_color from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette from supervision.geometry.core import Position -class ColorMap(Enum): - """ - Enum for annotator color mapping. - """ - - INDEX = "index" - CLASS = "class" - TRACK = "track" - - -def resolve_color_idx( - detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS -) -> int: - if detection_idx >= len(detections): - raise ValueError( - f"Detection index {detection_idx}" - f"is out of bounds for detections of length {len(detections)}" - ) - - if color_map == ColorMap.INDEX: - return detection_idx - elif color_map == ColorMap.CLASS: - if detections.class_id is None: - raise ValueError( - "Could not resolve color by class because" - "Detections do not have class_id" - ) - return detections.class_id[detection_idx] - elif color_map == ColorMap.TRACK: - if detections.tracker_id is None: - raise ValueError( - "Could not resolve color by track because" - "Detections do not have tracker_id" - ) - return detections.tracker_id[detection_idx] - - -def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color: - if isinstance(color, ColorPalette): - return color.by_idx(idx) - else: - return color - - -class BaseAnnotator(ABC): - @abstractmethod - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - pass - - class BoundingBoxAnnotator(BaseAnnotator): """ A class for drawing bounding boxes on an image using provided detections. @@ -350,81 +300,6 @@ class LabelAnnotator(BaseAnnotator): return scene -class BoxMaskAnnotator(BaseAnnotator): - def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - opacity: float = 0.5, - color_by_track: bool = False, - ): - self.color: Union[Color, ColorPalette] = color - self.opacity = opacity - self.color_by_track = color_by_track - - def annotate( - self, - scene: np.ndarray, - detections: Detections, - ) -> np.ndarray: - """ - Overlays the rectangle masks on the given image based on the - provided detections, with a specified opacity. - - Args: - scene (np.ndarray): The image on which the masks will be overlaid - detections (Detections): The detections for which the masks will be overlaid - Returns: - np.ndarray: The image with the masks overlaid - Example: - ```python - >>> import supervision as sv - - >>> classes = ['person', ...] - >>> image = ... - >>> detections = sv.Detections(...) - - >>> box_mask_annotator = sv.MaskAnnotator() - >>> annotated_frame = box_mask_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) - ``` - """ - overlay_img = np.zeros_like(scene, np.uint8) - 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 - ) - cv2.rectangle( - img=overlay_img, - pt1=(x1, y1), - pt2=(x2, y2), - color=color.as_bgr(), - thickness=-1, - ) - - mask = overlay_img.astype(bool) - scene[mask] = cv2.addWeighted( - scene, self.opacity, overlay_img, 1 - self.opacity, 0 - )[mask] - return scene - - def default_label_formatter(detections: Detections) -> List[str]: return [str(class_id) for class_id in detections.class_id] diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index c84ca826..83b6e0f8 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -147,64 +147,6 @@ class BoxAnnotator: return scene -class MaskAnnotator: - """ - A class for overlaying masks on an image using detections provided. - - Attributes: - color (Union[Color, ColorPalette]): The color to fill the mask, - can be a single color or a color palette - """ - - def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - ): - self.color: Union[Color, ColorPalette] = color - - def annotate( - self, scene: np.ndarray, detections: Detections, opacity: float = 0.5 - ) -> np.ndarray: - """ - Overlays the masks on the given image based on the provided detections, - with a specified opacity. - - Args: - scene (np.ndarray): The image on which the masks will be overlaid - detections (Detections): The detections for which the - masks will be overlaid - opacity (float): The opacity of the masks, between 0 and 1, default is 0.5 - - Returns: - np.ndarray: The image with the masks overlaid - """ - if detections.mask is None: - return scene - - for i in np.flip(np.argsort(detections.area)): - 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 - ) - - mask = detections.mask[i] - colored_mask = np.zeros_like(scene, dtype=np.uint8) - colored_mask[:] = color.as_bgr() - - scene = np.where( - np.expand_dims(mask, axis=-1), - np.uint8(opacity * colored_mask + (1 - opacity) * scene), - scene, - ) - - return scene - - class Trace: def __init__( self,