From 3e8906d83079e7520cd0197c2d4729d4836a40bd Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 28 Sep 2023 23:24:22 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=98=20docs=20updates=20for=20`LabelAnn?= =?UTF-8?q?otator`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/detection/annotate.md | 4 + supervision/__init__.py | 2 - supervision/annotators/core.py | 251 ++++++--------------------------- 3 files changed, 45 insertions(+), 212 deletions(-) diff --git a/docs/detection/annotate.md b/docs/detection/annotate.md index bf757157..8568139f 100644 --- a/docs/detection/annotate.md +++ b/docs/detection/annotate.md @@ -13,3 +13,7 @@ ## BoxCornerAnnotator :::supervision.annotators.core.BoxCornerAnnotator + +## LabelAnnotator + +:::supervision.annotators.core.LabelAnnotator \ No newline at end of file diff --git a/supervision/__init__.py b/supervision/__init__.py index ca7fad66..cce46b41 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -10,10 +10,8 @@ from supervision.annotators.core import ( BoundingBoxAnnotator, BoxCornerAnnotator, EllipseAnnotator, - LabelAdvancedAnnotator, LabelAnnotator, MaskAnnotator, - TraceAnnotator, ) from supervision.classification.core import Classifications from supervision.dataset.core import ( diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index a1a79726..5dd0871c 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -307,6 +307,9 @@ class BoxCornerAnnotator(BaseAnnotator): class LabelAnnotator: + """ + A class for annotating labels on an image using provided detections. + """ def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), @@ -317,6 +320,19 @@ class LabelAnnotator: text_position: Position = Position.TOP_LEFT, color_map: str = "class", ): + """ + Args: + color (Union[Color, ColorPalette]): The color or color palette to use for + annotating the text background. + text_color (Color): The color to use for the text. + text_scale (float): Font scale for the text. + text_thickness (int): Thickness of the text characters. + text_padding (int): Padding around the text within its background box. + text_position (Position): Position of the text relative to the detection. + Possible values are defined in the `Position` enum. + color_map (str): Strategy for mapping colors to annotations. + Options are `index`, `class`, or `track`. + """ self.color: Union[Color, ColorPalette] = color self.text_color: Color = text_color self.text_scale: float = text_scale @@ -367,6 +383,31 @@ class LabelAnnotator: detections: Detections, labels: List[str] = None, ) -> np.ndarray: + """ + Annotates the given scene with labels based on the provided detections. + + Args: + scene (np.ndarray): The image where labels will be drawn. + detections (Detections): Object detections to annotate. + labels (List[str]): Optional. Custom labels for each detection. + + Returns: + np.ndarray: The annotated image. + + Example: + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> label_annotator = sv.LabelAnnotator() + >>> annotated_frame = label_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + """ font = cv2.FONT_HERSHEY_SIMPLEX for detection_idx in range(len(detections)): detection_xyxy = detections.xyxy[detection_idx].astype(int) @@ -416,213 +457,3 @@ class LabelAnnotator: lineType=cv2.LINE_AA, ) return scene - - -class LabelAdvancedAnnotator(BaseAnnotator): - def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - text_color: Color = Color.black(), - text_padding: int = 20, - color_by_track: bool = False, - font: Optional[str] = None, - font_size: Optional[int] = 15, - ): - if font and os.path.exists(font): - self.font = ImageFont.truetype(font, font_size) - else: - self.font = ImageFont.load_default() - self.color: Union[Color, ColorPalette] = color - self.text_color: Color = text_color - self.text_padding: int = text_padding - self.color_by_track = color_by_track - - def annotate( - self, - scene: np.ndarray, - detections: Detections, - labels: Optional[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 - labels (Optional[List[str]]): An optional list of labels corresponding - to each detection. If `labels` are not provided, - corresponding `class_id` will be used as label. - Returns: - np.ndarray: The image with the bounding boxes drawn on it - - Example: - ```python - >>> import supervision as sv - - >>> classes = ['person', ...] - >>> image = ... - >>> detections = sv.Detections(...) - - >>> pil_label_annotator = sv.LabelAdvancedAnnotator() - >>> labels = [ - ... f"{classes[class_id]} {confidence:0.2f}" - ... for _, _, confidence, class_id, _ - ... in detections - ... ] - >>> annotated_frame = pil_label_annotator.annotate( - ... scene=image.copy(), - ... detections=detections, - ... labels=labels, - ... ) - ``` - """ - pil_image = Image.fromarray(scene) - draw = ImageDraw.Draw(pil_image) - text_color = "#fff" - - 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 - ) - - text = ( - f"{idx}" - if (labels is None or len(detections) != len(labels)) - else labels[i] - ) - - text_bbox = draw.textbbox((x1, y1), text, font=self.font) - - text_height = text_bbox[3] - text_bbox[1] - text_width = text_bbox[2] - text_bbox[0] - - text_x = x1 + self.text_padding / 2 - text_y = y1 - self.text_padding / 2 - text_height - - text_background_x1 = x1 - text_background_y1 = y1 - self.text_padding / 2 - text_height - - text_background_x2 = x1 + 2 * self.text_padding / 2 + text_width - text_background_y2 = y1 # correct - - draw.rectangle( - ( - text_background_x1, - text_background_y1, - text_background_x2, - text_background_y2, - ), - fill=color.as_bgr(), - ) - draw.text((text_x, text_y), text, font=self.font, fill=text_color) - - scene = np.asarray(pil_image) - return scene - - -class TraceAnnotator(BaseAnnotator): - """ - A class for drawing trajectory of a tracker on an image using detections provided. - - Attributes: - color (Union[Color, ColorPalette]): The color to draw the trajectory, - can be a single color or a color palette - color_by_track (bool): Whther to use tracker id to pick the color - position (Optional[Position]): Choose position of trajectory such as - center position, top left corner, etc - trace_length (int): Length of the previous points - thickness (int): thickness of the line - """ - - def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - color_by_track: bool = False, - position: Optional[Position] = Position.CENTER, - trace_length: int = 30, - thickness: int = 2, - ): - self.color: Union[Color, ColorPalette] = color - self.color_by_track = color_by_track - self.position = position - self.tracker_storage = defaultdict(lambda: []) - self.trace_length = trace_length - self.thickness = thickness - - def annotate( - self, scene: np.ndarray, detections: Detections, **kwargs - ) -> np.ndarray: - """ - Draw the object trajectory based on history of tracked objects - - Args: - scene (np.ndarray): The image on which the trace will be drawn - detections (Detections): The detections for trajectory and points - - Returns: - np.ndarray: The image with the masks overlaid - Example: - ```python - >>> import supervision as sv - - >>> classes = ['person', ...] - >>> image = ... - >>> detections = sv.Detections(...) - - >>> trace_annotator = sv.TraceAnnotator() - >>> annotated_frame = trace_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) - ``` - """ - if detections.tracker_id is None: - return scene - - anchor_points = detections.get_anchor_coordinates(anchor=self.position) - - for i, tracker_id in enumerate(detections.tracker_id): - track = self.tracker_storage[tracker_id] - track.append((anchor_points[i][0], anchor_points[i][1])) - if len(track) > self.trace_length: - track.pop(0) - points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2)) - - if self.color_by_track: - 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.polylines( - scene, - [points], - isClosed=False, - color=color.as_bgr(), - thickness=self.thickness, - ) - - return scene