Merge pull request #1786 from Armaggheddon/refactor/heatmap-perf

perf: improved hsv color mapping in HeatMapAnnotator by 20x
roboflow · Feb 19, 2025 · acf1ef6 · acf1ef6
2 parents 752818e + fefd005
commit acf1ef6
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py
@@ -1857,11 +1857,9 @@ def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
         temp = temp.astype(np.uint8)
         if self.kernel_size is not None:
             temp = cv2.blur(temp, (self.kernel_size, self.kernel_size))
-        hsv = np.zeros(scene.shape)
+        hsv = np.full(scene.shape, 255, dtype=np.uint8)
         hsv[..., 0] = temp
-        hsv[..., 1] = 255
-        hsv[..., 2] = 255
-        temp = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
+        temp = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
         mask = cv2.cvtColor(self.heat_mask.astype(np.uint8), cv2.COLOR_GRAY2BGR) > 0
         scene[mask] = cv2.addWeighted(temp, self.opacity, scene, 1 - self.opacity, 0)[
             mask