diff --git a/supervision/__init__.py b/supervision/__init__.py index e3fcd26d..fa6e6044 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -13,6 +13,7 @@ from supervision.annotators.core import ( BoxMaskAnnotator, CircleAnnotator, EllipseAnnotator, + HaloAnnotator, LabelAnnotator, MaskAnnotator, TraceAnnotator, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 2ee7c972..494c94ad 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -229,6 +229,88 @@ class BoxMaskAnnotator(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 = 0.8, + color_map: str = "class", + kernel_size: int = 40, + ): + """ + 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`. + 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: + """ + 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, (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) + return scene + + class EllipseAnnotator(BaseAnnotator): """ A class for drawing ellipses on an image using provided detections.