add params for hsv

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