From c5ff58910dd06fb03cbe5b1d3cdbcb1f7f4b2e19 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Tue, 18 Jun 2024 21:32:40 +0300 Subject: [PATCH] Reduce number of copy operations --- supervision/annotators/core.py | 55 ++++++++++++++---------------- supervision/detection/annotate.py | 4 +-- supervision/keypoint/annotators.py | 6 ++-- supervision/utils/conversion.py | 6 ++-- supervision/utils/image.py | 10 +++--- 5 files changed, 39 insertions(+), 42 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 6ea97a3b..38ae6a36 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -19,8 +19,8 @@ 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, - convert_for_rich_text_annotation, + ensure_cv2_image_for_annotation, + ensure_pil_image_for_annotation, ) from supervision.utils.image import crop_image, overlay_image, scale_image @@ -48,7 +48,7 @@ class BoundingBoxAnnotator(BaseAnnotator): self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -129,7 +129,7 @@ class OrientedBoxAnnotator(BaseAnnotator): self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -217,7 +217,7 @@ class MaskAnnotator(BaseAnnotator): self.opacity = opacity self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -273,10 +273,9 @@ class MaskAnnotator(BaseAnnotator): mask = detections.mask[detection_idx] colored_mask[mask] = color.as_bgr() - blended_scene = cv2.addWeighted( - colored_mask, self.opacity, scene, 1 - self.opacity, 0 + cv2.addWeighted( + colored_mask, self.opacity, scene, 1 - self.opacity, 0, dst=scene ) - np.copyto(scene, blended_scene.astype(np.uint8)) return scene @@ -307,7 +306,7 @@ class PolygonAnnotator(BaseAnnotator): self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -393,7 +392,7 @@ class ColorAnnotator(BaseAnnotator): self.color_lookup: ColorLookup = color_lookup self.opacity = opacity - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -432,7 +431,6 @@ class ColorAnnotator(BaseAnnotator): ![box-mask-annotator-example](https://media.roboflow.com/ supervision-annotator-examples/box-mask-annotator-example-purple.png) """ - mask_image = scene.copy() scene_with_boxes = scene.copy() for detection_idx in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int) @@ -452,10 +450,9 @@ class ColorAnnotator(BaseAnnotator): thickness=-1, ) - blended_scene = cv2.addWeighted( - scene_with_boxes, self.opacity, mask_image, 1 - self.opacity, gamma=0 + cv2.addWeighted( + scene_with_boxes, self.opacity, scene, 1 - self.opacity, gamma=0, dst=scene ) - np.copyto(scene, blended_scene) return scene @@ -490,7 +487,7 @@ class HaloAnnotator(BaseAnnotator): self.color_lookup: ColorLookup = color_lookup self.kernel_size: int = kernel_size - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -589,7 +586,7 @@ class EllipseAnnotator(BaseAnnotator): self.end_angle: int = end_angle self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -680,7 +677,7 @@ class BoxCornerAnnotator(BaseAnnotator): self.corner_length: int = corner_length self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -768,7 +765,7 @@ class CircleAnnotator(BaseAnnotator): self.thickness: int = thickness self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -858,7 +855,7 @@ class DotAnnotator(BaseAnnotator): self.position: Position = position self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -952,7 +949,7 @@ class LabelAnnotator(BaseAnnotator): self.text_anchor: Position = text_position self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -1158,7 +1155,7 @@ class RichLabelAnnotator(BaseAnnotator): else: self.font = self._load_default_font(font_size) - @convert_for_rich_text_annotation + @ensure_pil_image_for_annotation def annotate( self, scene: ImageType, @@ -1289,7 +1286,7 @@ class BlurAnnotator(BaseAnnotator): """ self.kernel_size: int = kernel_size - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -1374,7 +1371,7 @@ class TraceAnnotator(BaseAnnotator): self.thickness = thickness self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -1481,7 +1478,7 @@ class HeatMapAnnotator(BaseAnnotator): self.top_hue = top_hue self.low_hue = low_hue - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate(self, scene: ImageType, detections: Detections) -> ImageType: """ Annotates the scene with a heatmap based on the provided detections. @@ -1557,7 +1554,7 @@ class PixelateAnnotator(BaseAnnotator): """ self.pixel_size: int = pixel_size - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -1645,7 +1642,7 @@ class TriangleAnnotator(BaseAnnotator): self.position: Position = position self.color_lookup: ColorLookup = color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -1741,7 +1738,7 @@ class RoundBoxAnnotator(BaseAnnotator): raise ValueError("roundness attribute must be float between (0, 1.0]") self.roundness: float = roundness - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -1878,7 +1875,7 @@ class PercentageBarAnnotator(BaseAnnotator): if border_thickness is None: self.border_thickness = int(0.15 * self.height) - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, @@ -2051,7 +2048,7 @@ class CropAnnotator(BaseAnnotator): self.border_thickness: int = border_thickness self.border_color_lookup: ColorLookup = border_color_lookup - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index f496b248..eb52171c 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -5,7 +5,7 @@ import cv2 from supervision.annotators.base import ImageType from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette -from supervision.utils.conversion import convert_for_annotation_method +from supervision.utils.conversion import ensure_cv2_image_for_annotation from supervision.utils.internal import deprecated @@ -46,7 +46,7 @@ class BoxAnnotator: "`BoxAnnotator` is deprecated and will be removed in " "`supervision-0.22.0`. Use `BoundingBoxAnnotator` and `LabelAnnotator` instead" ) - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate( self, scene: ImageType, diff --git a/supervision/keypoint/annotators.py b/supervision/keypoint/annotators.py index e6ff1fcf..256b1b62 100644 --- a/supervision/keypoint/annotators.py +++ b/supervision/keypoint/annotators.py @@ -11,7 +11,7 @@ from supervision.draw.color import Color from supervision.draw.utils import draw_rounded_rectangle from supervision.keypoint.core import KeyPoints from supervision.keypoint.skeletons import SKELETONS_BY_VERTEX_COUNT -from supervision.utils.conversion import convert_for_annotation_method +from supervision.utils.conversion import ensure_cv2_image_for_annotation class BaseKeyPointAnnotator(ABC): @@ -41,7 +41,7 @@ class VertexAnnotator(BaseKeyPointAnnotator): self.color = color self.radius = radius - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType: """ Annotates the given scene with skeleton vertices based on the provided key @@ -117,7 +117,7 @@ class EdgeAnnotator(BaseKeyPointAnnotator): self.thickness = thickness self.edges = edges - @convert_for_annotation_method + @ensure_cv2_image_for_annotation def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType: """ Annotates the given scene by drawing lines between specified key points to form diff --git a/supervision/utils/conversion.py b/supervision/utils/conversion.py index 1a73a1b4..646d7dcf 100644 --- a/supervision/utils/conversion.py +++ b/supervision/utils/conversion.py @@ -8,7 +8,7 @@ from PIL import Image from supervision.annotators.base import ImageType -def convert_for_annotation_method(annotate_func): +def ensure_cv2_image_for_annotation(annotate_func): """ Decorates `BaseAnnotator.annotate` implementations, converts scene to an image type used internally by the annotators, converts back when annotation @@ -33,7 +33,7 @@ def convert_for_annotation_method(annotate_func): return wrapper -def convert_for_image_processing(image_processing_fun): +def ensure_cv2_image_for_processing(image_processing_fun): """ Decorates image processing functions that accept np.ndarray, converting `image` to np.ndarray, converts back when processing is complete. @@ -56,7 +56,7 @@ def convert_for_image_processing(image_processing_fun): return wrapper -def convert_for_rich_text_annotation(annotate_func): +def ensure_pil_image_for_annotation(annotate_func): """ Decorates image processing functions that accept np.ndarray, converting `image` to PIL image, converts back when processing is complete. diff --git a/supervision/utils/image.py b/supervision/utils/image.py index 9f2e1783..d6972189 100644 --- a/supervision/utils/image.py +++ b/supervision/utils/image.py @@ -14,8 +14,8 @@ from supervision.draw.color import Color, unify_to_bgr from supervision.draw.utils import calculate_optimal_text_scale, draw_text from supervision.geometry.core import Point from supervision.utils.conversion import ( - convert_for_image_processing, cv2_to_pillow, + ensure_cv2_image_for_processing, images_to_cv2, ) from supervision.utils.iterables import create_batches, fill @@ -25,7 +25,7 @@ RelativePosition = Literal["top", "bottom"] MAX_COLUMNS_FOR_SINGLE_ROW_GRID = 3 -@convert_for_image_processing +@ensure_cv2_image_for_processing def crop_image( image: ImageType, xyxy: Union[npt.NDArray[int], List[int], Tuple[int, int, int, int]], @@ -86,7 +86,7 @@ def crop_image( return image[y_min:y_max, x_min:x_max] -@convert_for_image_processing +@ensure_cv2_image_for_processing def scale_image(image: ImageType, scale_factor: float) -> ImageType: """ Scales the given image based on the given scale factor. @@ -143,7 +143,7 @@ def scale_image(image: ImageType, scale_factor: float) -> ImageType: return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR) -@convert_for_image_processing +@ensure_cv2_image_for_processing def resize_image( image: ImageType, resolution_wh: Tuple[int, int], @@ -216,7 +216,7 @@ def resize_image( return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR) -@convert_for_image_processing +@ensure_cv2_image_for_processing def letterbox_image( image: ImageType, resolution_wh: Tuple[int, int],