From e3cb41712924d0ddbd09303ceb7203319082a083 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Tue, 11 Jun 2024 16:15:23 +0300 Subject: [PATCH] RichLabelAnnotator fix: convert back to numpy --- supervision/annotators/core.py | 32 +++++++++++++++----------------- supervision/utils/conversion.py | 21 +++++++++++++++++++++ 2 files changed, 36 insertions(+), 17 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 1c0ad8af..4f1e6760 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -3,7 +3,7 @@ from typing import List, Optional, Tuple, Union import cv2 import numpy as np -from PIL import Image, ImageDraw, ImageFont +from PIL import ImageDraw, ImageFont from supervision.annotators.base import BaseAnnotator, ImageType from supervision.annotators.utils import ( @@ -18,7 +18,10 @@ from supervision.detection.utils import clip_boxes, mask_to_polygons from supervision.draw.color import Color, ColorPalette from supervision.draw.utils import draw_polygon from supervision.geometry.core import Position -from supervision.utils.conversion import convert_for_annotation_method +from supervision.utils.conversion import ( + convert_for_annotation_method, + convert_for_rich_text_annotation, +) from supervision.utils.image import crop_image, overlay_image, scale_image @@ -947,7 +950,7 @@ class LabelAnnotator: self, scene: ImageType, detections: Detections, - labels: List[str] = None, + labels: Optional[List[str]] = None, custom_color_lookup: Optional[np.ndarray] = None, ) -> ImageType: """ @@ -958,7 +961,7 @@ class LabelAnnotator: `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. + labels (Optional[List[str]]): Custom labels for each detection. custom_color_lookup (Optional[np.ndarray]): Custom color lookup array. Allows to override the default color mapping strategy. @@ -968,7 +971,7 @@ class LabelAnnotator: Example: ```python - import supervision as sv + import supervision as sv image = ... detections = sv.Detections(...) @@ -996,12 +999,9 @@ class LabelAnnotator: ).astype(int) if labels is not None and len(labels) != len(detections): raise ValueError( - f"The number of labels provided ({len(labels)}) does not match the " - f"number of detections ({len(detections)}). Each detection should have " - f"a corresponding label. This discrepancy can occur if the labels and " - f"detections are not aligned or if an incorrect number of labels has " - f"been provided. Please ensure that the labels array has the same " - f"length as the Detections object." + f"The number of labels ({len(labels)}) does not match the " + f"number of detections ({len(detections)}). Each detection " + f"should have exactly 1 label." ) for detection_idx, center_coordinates in enumerate(anchors_coordinates): @@ -1113,7 +1113,7 @@ class RichLabelAnnotator: self, color: Union[Color, ColorPalette] = ColorPalette.DEFAULT, text_color: Color = Color.WHITE, - font_path: str = None, + font_path: Optional[str] = None, font_size: int = 10, text_padding: int = 10, text_position: Position = Position.TOP_LEFT, @@ -1125,8 +1125,8 @@ class RichLabelAnnotator: 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. - font_path (str): Path to the font file (e.g., ".ttf" or ".otf") to use for - rendering text. If `None`, the default PIL font will be used. + font_path (Optional[str]): Path to the font file (e.g., ".ttf" or ".otf") + to use for rendering text. If `None`, the default PIL font will be used. font_size (int): Font size for the text. text_padding (int): Padding around the text within its background box. text_position (Position): Position of the text relative to the detection. @@ -1151,6 +1151,7 @@ class RichLabelAnnotator: else: self.font = ImageFont.load_default(size=font_size) + @convert_for_rich_text_annotation def annotate( self, scene: ImageType, @@ -1197,8 +1198,6 @@ class RichLabelAnnotator: ``` """ - if isinstance(scene, np.ndarray): - scene = Image.fromarray(cv2.cvtColor(scene, cv2.COLOR_BGR2RGB)) draw = ImageDraw.Draw(scene) anchors_coordinates = detections.get_anchors_coordinates( anchor=self.text_anchor @@ -1258,7 +1257,6 @@ class RichLabelAnnotator: font=self.font, fill=self.text_color.as_rgb(), ) - return scene diff --git a/supervision/utils/conversion.py b/supervision/utils/conversion.py index 8ddce969..b3c63a86 100644 --- a/supervision/utils/conversion.py +++ b/supervision/utils/conversion.py @@ -51,6 +51,27 @@ def convert_for_image_processing(image_processing_fun): return wrapper +def convert_for_rich_text_annotation(annotate_func): + """ + Decorates image processing functions that accept np.ndarray, converting `image` to + PIL image, converts back when processing is complete. + """ + + @wraps(annotate_func) + def wrapper(self, scene: ImageType, *args, **kwargs): + if isinstance(scene, np.ndarray): + scene = cv2_to_pillow(scene) + annotated = annotate_func(self, scene, *args, **kwargs) + return pillow_to_cv2(image=annotated) + + if isinstance(scene, Image.Image): + return annotate_func(self, scene, *args, **kwargs) + + raise ValueError(f"Unsupported image type: {type(scene)}") + + return wrapper + + def images_to_cv2(images: List[ImageType]) -> List[np.ndarray]: """ Converts images provided either as Pillow images or OpenCV