From fefd005c2c746fe023d8fd4c77fa3367f4738b33 Mon Sep 17 00:00:00 2001 From: Armaggheddon Date: Mon, 10 Feb 2025 20:59:45 +0100 Subject: [PATCH] perf: improved hsv color mapping in HeatMapAnnotator by 20x --- supervision/annotators/core.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index a2b2ed56..c7c41b03 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1944,11 +1944,9 @@ class HeatMapAnnotator(BaseAnnotator): 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