diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 072e1026..a9aa0da6 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -97,6 +97,8 @@ class HeatmapAnnotator: opacity: float = 0.2, radius: int = 40, kernel_size: Optional[int] = 25, + top_hue: int = 0, #red + low_hue: int = 125, #blue ): """ Args: @@ -105,12 +107,19 @@ class HeatmapAnnotator: opacity (float): Opacity of the overlay mask. Must be between `0` and `1`. radius (int): Radius of the heat circle. kernel_size (Optional[int]): Kernel size for blurring the heatmap. + top_hue (int): Hue of the top of the heatmap. + Defaults to `0` (red). + low_hue (int): Hue of the bottom of the heatmap. + Defaults to `125` (blue). """ self.position = position self.opacity = opacity self.radius = radius self.kernel_size = kernel_size self.heatmask = None + self.top_hue = top_hue + self.low_hue = low_hue + def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: """ @@ -156,19 +165,17 @@ class HeatmapAnnotator: cv2.circle(mask, (int(xy[0]), int(xy[1])), self.radius, 1, -1) self.heatmask = mask + self.heatmask temp = self.heatmask.copy() - temp = 100 - temp / temp.max() * 90 + temp = self.low_hue-temp/temp.max()*(self.low_hue-self.top_hue) temp = temp.astype(np.uint8) if self.kernel_size is not None: - temp = cv2.GaussianBlur(temp, (self.kernel_size, self.kernel_size), 0) + temp = cv2.blur(temp, (self.kernel_size, self.kernel_size)) hsv = np.zeros(scene.shape) hsv[..., 0] = temp hsv[..., 1] = 255 hsv[..., 2] = 255 - mask = temp > 0 temp = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) - scene[mask] = cv2.addWeighted(temp, self.opacity, scene, 1 - self.opacity, 0)[ - mask - ] + mask = cv2.cvtColor(self.heatmask.astype(np.uint8), cv2.COLOR_GRAY2BGR)>0 + scene[mask] = cv2.addWeighted( temp, self.opacity,scene, 1-self.opacity, 0)[mask] return scene