From 72dc07e33245f66d5f85681a14a16c8cde756efd Mon Sep 17 00:00:00 2001 From: kapter Date: Sat, 7 Oct 2023 01:11:49 +0300 Subject: [PATCH 1/4] added HaloAnnotator --- supervision/__init__.py | 1 + supervision/annotators/core.py | 79 ++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/supervision/__init__.py b/supervision/__init__.py index d6492af1..4c2d3dd1 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -15,6 +15,7 @@ from supervision.annotators.core import ( LabelAnnotator, MaskAnnotator, TraceAnnotator, + HaloAnnotator, ) from supervision.classification.core import Classifications from supervision.dataset.core import ( diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 8edca92b..815b52c6 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -157,6 +157,85 @@ class MaskAnnotator(BaseAnnotator): ) return scene +class HaloAnnotator(BaseAnnotator): + """ + A class for drawing Halos on an image using provided detections. + """ + + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + opacity: float = .8, + color_map: str = "class", + ): + """ + Args: + color (Union[Color, ColorPalette]): The color or color palette to use for + annotating detections. + opacity (float): Opacity of the overlay mask. Must be between `0` and `1`. + color_map (str): Strategy for mapping colors to annotations. + Options are `index`, `class`, or `track`. + """ + self.color: Union[Color, ColorPalette] = color + self.opacity = opacity + self.color_map: ColorMap = ColorMap(color_map) + + def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: + """ + Annotates the given scene with halos based on the provided detections. + + Args: + scene (np.ndarray): The image where masks will be drawn. + detections (Detections): Object detections to annotate. + + Returns: + np.ndarray: The annotated image. + + Example: + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> halo_annotator = sv.HaloAnnotator() + >>> annotated_frame = halo_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + + ![halo-annotator-example](https://media.roboflow.com/ + supervision-annotator-examples/halo-annotator-example.png) + """ + if detections.mask is None: + return scene + colored_mask = np.zeros_like(scene, dtype=np.uint8) + fmask = np.array( + [False]*scene.shape[0]*scene.shape[1] + ).reshape(scene.shape[0],scene.shape[1]) + + + for detection_idx in np.flip(np.argsort(detections.area)): + idx = resolve_color_idx( + detections=detections, + detection_idx=detection_idx, + color_map=self.color_map, + ) + color = resolve_color(color=self.color, idx=idx) + mask = detections.mask[detection_idx] + fmask = np.logical_or(fmask,mask) + color_bgr = color.as_bgr() + colored_mask[mask] = color_bgr + colored_mask = cv2.blur(colored_mask,(20,20)) + colored_mask[fmask] = [0,0,0] + gray = cv2.cvtColor(colored_mask, cv2.COLOR_BGR2GRAY) + _, tresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY) + mask = tresh>0 + scene[mask] = cv2.addWeighted(colored_mask, self.opacity, scene, 1, 0)[mask] + + + return scene class EllipseAnnotator(BaseAnnotator): """ From 44e4eeaa69ad9524826187850ae2d508e537f3dc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 6 Oct 2023 22:33:32 +0000 Subject: [PATCH 2/4] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/__init__.py | 2 +- supervision/annotators/core.py | 20 ++++++++++---------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 4c2d3dd1..075c7e97 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -12,10 +12,10 @@ from supervision.annotators.core import ( BoxCornerAnnotator, CircleAnnotator, EllipseAnnotator, + HaloAnnotator, LabelAnnotator, MaskAnnotator, TraceAnnotator, - HaloAnnotator, ) from supervision.classification.core import Classifications from supervision.dataset.core import ( diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 815b52c6..a0e2da26 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -157,6 +157,7 @@ class MaskAnnotator(BaseAnnotator): ) return scene + class HaloAnnotator(BaseAnnotator): """ A class for drawing Halos on an image using provided detections. @@ -165,7 +166,7 @@ class HaloAnnotator(BaseAnnotator): def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), - opacity: float = .8, + opacity: float = 0.8, color_map: str = "class", ): """ @@ -211,10 +212,9 @@ class HaloAnnotator(BaseAnnotator): if detections.mask is None: return scene colored_mask = np.zeros_like(scene, dtype=np.uint8) - fmask = np.array( - [False]*scene.shape[0]*scene.shape[1] - ).reshape(scene.shape[0],scene.shape[1]) - + fmask = np.array([False] * scene.shape[0] * scene.shape[1]).reshape( + scene.shape[0], scene.shape[1] + ) for detection_idx in np.flip(np.argsort(detections.area)): idx = resolve_color_idx( @@ -224,19 +224,19 @@ class HaloAnnotator(BaseAnnotator): ) color = resolve_color(color=self.color, idx=idx) mask = detections.mask[detection_idx] - fmask = np.logical_or(fmask,mask) + fmask = np.logical_or(fmask, mask) color_bgr = color.as_bgr() colored_mask[mask] = color_bgr - colored_mask = cv2.blur(colored_mask,(20,20)) - colored_mask[fmask] = [0,0,0] + colored_mask = cv2.blur(colored_mask, (20, 20)) + colored_mask[fmask] = [0, 0, 0] gray = cv2.cvtColor(colored_mask, cv2.COLOR_BGR2GRAY) _, tresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY) - mask = tresh>0 + mask = tresh > 0 scene[mask] = cv2.addWeighted(colored_mask, self.opacity, scene, 1, 0)[mask] - return scene + class EllipseAnnotator(BaseAnnotator): """ A class for drawing ellipses on an image using provided detections. From 89c9f88aeb8f9485effb79e58bde267aa3a49206 Mon Sep 17 00:00:00 2001 From: kapter Date: Sat, 7 Oct 2023 23:12:11 +0300 Subject: [PATCH 3/4] add kernel size var and change method of overlay --- supervision/annotators/core.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index a0e2da26..3a5e86e2 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -168,6 +168,7 @@ class HaloAnnotator(BaseAnnotator): color: Union[Color, ColorPalette] = ColorPalette.default(), opacity: float = 0.8, color_map: str = "class", + kernel_size: int = 40, ): """ Args: @@ -176,10 +177,12 @@ class HaloAnnotator(BaseAnnotator): opacity (float): Opacity of the overlay mask. Must be between `0` and `1`. color_map (str): Strategy for mapping colors to annotations. Options are `index`, `class`, or `track`. + kernel_size (int): The size of the average pooling kernel used for creating the halo. """ self.color: Union[Color, ColorPalette] = color self.opacity = opacity self.color_map: ColorMap = ColorMap(color_map) + self.kernel_size: int = kernel_size def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: """ @@ -212,10 +215,10 @@ class HaloAnnotator(BaseAnnotator): if detections.mask is None: return scene colored_mask = np.zeros_like(scene, dtype=np.uint8) - fmask = np.array([False] * scene.shape[0] * scene.shape[1]).reshape( - scene.shape[0], scene.shape[1] - ) - + fmask = np.array( + [False]*scene.shape[0]*scene.shape[1] + ).reshape(scene.shape[0],scene.shape[1]) + for detection_idx in np.flip(np.argsort(detections.area)): idx = resolve_color_idx( detections=detections, @@ -227,13 +230,13 @@ class HaloAnnotator(BaseAnnotator): fmask = np.logical_or(fmask, mask) color_bgr = color.as_bgr() colored_mask[mask] = color_bgr - colored_mask = cv2.blur(colored_mask, (20, 20)) + + colored_mask = cv2.blur(colored_mask, (self.kernel_size, self.kernel_size)) colored_mask[fmask] = [0, 0, 0] gray = cv2.cvtColor(colored_mask, cv2.COLOR_BGR2GRAY) - _, tresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY) - mask = tresh > 0 - scene[mask] = cv2.addWeighted(colored_mask, self.opacity, scene, 1, 0)[mask] - + alpha = self.opacity*gray/gray.max() + alpha_mask = alpha[:,:,np.newaxis] + scene= np.uint8(scene*(1-alpha_mask)+colored_mask*self.opacity) return scene From a83d3b4c7e454b61843cebca1fe333990903fa36 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Oct 2023 20:12:29 +0000 Subject: [PATCH 4/4] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/annotators/core.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 3a5e86e2..e2aa0db7 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -215,10 +215,10 @@ class HaloAnnotator(BaseAnnotator): if detections.mask is None: return scene colored_mask = np.zeros_like(scene, dtype=np.uint8) - fmask = np.array( - [False]*scene.shape[0]*scene.shape[1] - ).reshape(scene.shape[0],scene.shape[1]) - + fmask = np.array([False] * scene.shape[0] * scene.shape[1]).reshape( + scene.shape[0], scene.shape[1] + ) + for detection_idx in np.flip(np.argsort(detections.area)): idx = resolve_color_idx( detections=detections, @@ -234,9 +234,9 @@ class HaloAnnotator(BaseAnnotator): colored_mask = cv2.blur(colored_mask, (self.kernel_size, self.kernel_size)) colored_mask[fmask] = [0, 0, 0] gray = cv2.cvtColor(colored_mask, cv2.COLOR_BGR2GRAY) - alpha = self.opacity*gray/gray.max() - alpha_mask = alpha[:,:,np.newaxis] - scene= np.uint8(scene*(1-alpha_mask)+colored_mask*self.opacity) + alpha = self.opacity * gray / gray.max() + alpha_mask = alpha[:, :, np.newaxis] + scene = np.uint8(scene * (1 - alpha_mask) + colored_mask * self.opacity) return scene