diff --git a/supervision/annotators/composable.py b/supervision/annotators/composable.py deleted file mode 100644 index 9a7482d0..00000000 --- a/supervision/annotators/composable.py +++ /dev/null @@ -1,100 +0,0 @@ -from abc import ABC -from typing import List, Optional - -import numpy as np - -from supervision.annotators.core import ( - BoundingBoxAnnotator, - LabelAnnotator, - MaskAnnotator, -) -from supervision.detection.core import Detections - - -class ComposableAnnotator(ABC): - def __init__( - self, - ): - self.annotators = [] - - def annotate( - self, - scene: np.ndarray, - detections: Detections, - labels: Optional[List[str]] = None, - ) -> np.ndarray: - for annotator in self.annotators: - if isinstance(annotator, LabelAnnotator): - scene = annotator.annotate( - scene=scene, - detections=detections, - labels=labels, - ) - else: - scene = annotator.annotate(scene=scene, detections=detections) - return scene - - -class DetectionAnnotator(ComposableAnnotator): - """ - Highlevel API for drawing Object Detection output. - This will use Box and Label Annotators - Example: - ```python - >>> import supervision as sv - - >>> classes = ['person', ...] - >>> image = ... - >>> detections = sv.Detections(...) - - >>> detection_annotator = sv.DetectionAnnotator() - >>> annotated_frame = detection_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) - ``` - """ - - def __init__( - self, - color_by_track: bool = False, - skip_label: bool = False, - ): - super().__init__() - self.annotators = [BoundingBoxAnnotator(color_by_track=color_by_track)] - if not skip_label: - self.annotators.append(LabelAnnotator(color_by_track=color_by_track)) - - -class SegmentationAnnotator(ComposableAnnotator): - """ - High level API for drawing segmentation mask, bounding box and labels on an - image using provided detections. - Example: - ```python - >>> import supervision as sv - - >>> classes = ['person', ...] - >>> image = ... - >>> detections = sv.Detections(...) - - >>> segmentation_annotator = sv.SegmentationAnnotator() - >>> annotated_frame = segmentation_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) - ``` - """ - - def __init__( - self, - color_by_track: bool = False, - skip_label: bool = False, - ): - super().__init__() - self.annotators = [ - BoundingBoxAnnotator(color_by_track=color_by_track), - MaskAnnotator(color_by_track=color_by_track), - ] - if not skip_label: - self.annotators.append(LabelAnnotator(color_by_track=color_by_track)) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 449c24c2..0900add0 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,6 +1,6 @@ import os.path from collections import defaultdict -from typing import Callable, List, Optional, Union +from typing import List, Optional, Union, Tuple import cv2 import numpy as np @@ -162,25 +162,25 @@ class EllipseAnnotator(BaseAnnotator): self, color: Union[Color, ColorPalette] = ColorPalette.default(), thickness: int = 2, - color_map: str = "class", start_angle: int = -45, end_angle: int = 235, + color_map: str = "class", ): """ Args: color (Union[Color, ColorPalette]): The color or color palette to use for annotating detections. thickness (int): Thickness of the ellipse lines. - color_map (str): Strategy for mapping colors to annotations. - Options are `index`, `class`, or `track`. start_angle (int): Starting angle of the ellipse. end_angle (int): Ending angle of the ellipse. + color_map (str): Strategy for mapping colors to annotations. + Options are `index`, `class`, or `track`. """ self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness + self.start_angle: int = start_angle + self.end_angle: int = end_angle self.color_map: ColorMap = ColorMap(color_map) - self.start_angle = start_angle - self.end_angle = end_angle def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: """ @@ -241,22 +241,22 @@ class BoxCornerAnnotator(BaseAnnotator): self, color: Union[Color, ColorPalette] = ColorPalette.default(), thickness: int = 4, - color_map: str = "class", corner_length: int = 25, + color_map: str = "class", ): """ Args: color (Union[Color, ColorPalette]): The color or color palette to use for annotating detections. thickness (int): Thickness of the corner lines. + corner_length (int): Length of each corner line. color_map (str): Strategy for mapping colors to annotations. Options are `index`, `class`, or `track`. - corner_length (int): Length of each corner line. """ self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness + self.corner_length: int = corner_length self.color_map: ColorMap = ColorMap(color_map) - self.corner_length = corner_length def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: """ @@ -311,28 +311,7 @@ class BoxCornerAnnotator(BaseAnnotator): return scene -def default_label_formatter( - detections: Detections, -) -> List[str]: - return [str(class_id) for class_id in detections.class_id] - - -def build_label_formatter(classes: List[str]) -> Callable[[Detections], List[str]]: - def default_label_formatter(detections: Detections) -> List[str]: - return [classes[class_id] for class_id in detections.class_id] - - return default_label_formatter - - -class LabelAnnotator(BaseAnnotator): - """ - A class for putting text on an image using provided detections. - - Attributes: - color (Union[Color, ColorPalette]): The color to text on the image, - can be a single color or a color palette - """ - +class LabelAnnotator: def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), @@ -340,113 +319,70 @@ class LabelAnnotator(BaseAnnotator): text_scale: float = 0.5, text_thickness: int = 1, text_padding: int = 10, - color_by_track: bool = False, - classes: Optional[List[str]] = None, - label_formatter: Optional[ - Callable[[Detections], List[str]] - ] = default_label_formatter, + text_position: Position = Position.TOP_LEFT, + color_map: str = "class", ): - """ - Args: - color_by_track: pick color by tracker id - classes: Optional list of class name - label_formatter: Optional callback function for label generator, - avoided if classes is provided - """ self.color: Union[Color, ColorPalette] = color self.text_color: Color = text_color self.text_scale: float = text_scale self.text_thickness: int = text_thickness self.text_padding: int = text_padding - self.color_by_track = color_by_track - self.classes = classes - self.label_formatter = label_formatter + self.text_position: Position = text_position + self.color_map: ColorMap = ColorMap(color_map) + + @staticmethod + def resolve_text_background_xyxy( + detection_xyxy: Tuple[int, int, int, int], + text_wh: Tuple[int, int], + text_padding: int, + position: Position, + ) -> Tuple[int, int, int, int]: + padded_text_wh = (text_wh[0] + 2 * text_padding, text_wh[1] + 2 * text_padding) + x1, y1, x2, y2 = detection_xyxy + if position == Position.TOP_LEFT: + return x1, y1 - padded_text_wh[1], x1 + padded_text_wh[0], y1 def annotate( self, scene: np.ndarray, detections: Detections, + labels: List[str] = None, ) -> np.ndarray: - """ - Draws text on the frame using the detections provided and label. - - Args: - scene (np.ndarray): The image on which the bounding boxes will be drawn - detections (Detections): The detections for which - the bounding boxes will be drawn - Returns: - np.ndarray: The image with the bounding boxes drawn on it - - Example: - ```python - >>> import supervision as sv - - >>> classes = ['person', ...] - >>> image = ... - >>> detections = sv.Detections(...) - >>> label_annotator = sv.LabelAnnotator(classes=classes) - >>> annotated_frame = label_annotator.annotate( - ... scene=image.copy(), - ... detections=detections, - ... ) - ``` - """ font = cv2.FONT_HERSHEY_SIMPLEX - - labels = None - if self.classes: - label_formatter = build_label_formatter(self.classes) - labels = label_formatter(detections) - else: - labels = self.label_formatter(detections=detections) - - 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 + for detection_idx in range(len(detections)): + detection_xyxy = 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) text = ( - f"{idx}" + f"{detections.class_id[detection_idx]}" if (labels is None or len(detections) != len(labels)) - else labels[i] + else labels[detection_idx] ) - - text_width, text_height = cv2.getTextSize( + text_wh = cv2.getTextSize( text=text, fontFace=font, fontScale=self.text_scale, thickness=self.text_thickness, )[0] - text_x = x1 + self.text_padding - text_y = y1 - self.text_padding + text_background_xyxy = self.resolve_text_background_xyxy( + detection_xyxy=detection_xyxy, + text_wh=text_wh, + text_padding=self.text_padding, + position=self.text_position, + ) - text_background_x1 = x1 - text_background_y1 = y1 - 2 * self.text_padding - text_height - - text_background_x2 = x1 + 2 * self.text_padding + text_width - text_background_y2 = y1 + text_x = text_background_xyxy[0] + self.text_padding + text_y = text_background_xyxy[1] + self.text_padding + text_wh[1] cv2.rectangle( img=scene, - pt1=(text_background_x1, text_background_y1), - pt2=(text_background_x2, text_background_y2), + pt1=(text_background_xyxy[0], text_background_xyxy[1]), + pt2=(text_background_xyxy[2], text_background_xyxy[3]), color=color.as_bgr(), thickness=cv2.FILLED, ) @@ -463,10 +399,6 @@ class LabelAnnotator(BaseAnnotator): return scene -def default_label_formatter(detections: Detections) -> List[str]: - return [str(class_id) for class_id in detections.class_id] - - class LabelAdvancedAnnotator(BaseAnnotator): def __init__( self, diff --git a/test/annotators/test_core.py b/test/annotators/test_base.py similarity index 100% rename from test/annotators/test_core.py rename to test/annotators/test_base.py