diff --git a/src/supervision/annotators/core.py b/src/supervision/annotators/core.py index d754ac79..9c8dc1b9 100644 --- a/src/supervision/annotators/core.py +++ b/src/supervision/annotators/core.py @@ -2095,7 +2095,7 @@ class HeatMapAnnotator(BaseAnnotator): position: Position = Position.BOTTOM_CENTER, opacity: float = 0.2, radius: int = 40, - kernel_size: int = 25, + kernel_size: int | None = 25, top_hue: int = 0, low_hue: int = 125, ): @@ -2105,7 +2105,8 @@ class HeatMapAnnotator(BaseAnnotator): `BOTTOM_CENTER`. opacity: Opacity of the overlay mask, between 0 and 1. radius: Radius of the heat circle. - kernel_size: Kernel size for blurring the heatmap. + kernel_size: Kernel size for blurring the heatmap. Pass `None` + to disable blurring entirely. top_hue: Hue at the top of the heatmap. Defaults to 0 (red). low_hue: Hue at the bottom of the heatmap. Defaults to 125 (blue). """ @@ -2132,6 +2133,10 @@ class HeatMapAnnotator(BaseAnnotator): The annotated image, matching the type of `scene` (`numpy.ndarray` or `PIL.Image.Image`) + Note: + When `detections` is empty or no heat has accumulated yet, the + scene is returned unchanged without raising a ``RuntimeWarning``. + Example: ```python import supervision as sv @@ -2174,7 +2179,9 @@ class HeatMapAnnotator(BaseAnnotator): ) self.heat_mask = mask + self.heat_mask temp = self.heat_mask.copy() - temp = self.low_hue - temp / temp.max() * (self.low_hue - self.top_hue) + max_val = temp.max() + if max_val > 0: + temp = self.low_hue - temp / max_val * (self.low_hue - self.top_hue) temp = temp.astype(np.uint8) if self.kernel_size is not None: temp = cv2.blur(temp, (self.kernel_size, self.kernel_size)) diff --git a/tests/annotators/test_core.py b/tests/annotators/test_core.py index 5702b8f5..51428edb 100644 --- a/tests/annotators/test_core.py +++ b/tests/annotators/test_core.py @@ -2,6 +2,8 @@ Tests for supervision/annotators/core.py """ +import warnings + import numpy as np import pytest @@ -17,6 +19,7 @@ from supervision.annotators.core import ( DotAnnotator, EllipseAnnotator, HaloAnnotator, + HeatMapAnnotator, LabelAnnotator, MaskAnnotator, OrientedBoxAnnotator, @@ -361,6 +364,49 @@ class TestHaloAnnotator: assert np.array_equal(result_bool, result_uint8) +class TestHeatMapAnnotator: + """Tests for HeatMapAnnotator class""" + + def test_annotate_with_no_detections_does_not_warn( + self, test_image: np.ndarray + ) -> None: + """Empty detections must not trigger a divide-by-zero RuntimeWarning.""" + detections = Detections.empty() + annotator = HeatMapAnnotator() + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + result = annotator.annotate(scene=test_image.copy(), detections=detections) + assert np.array_equal(test_image, result) + + def test_annotate_with_single_detection(self, test_image: np.ndarray) -> None: + """Single detection must produce visible heat — result differs from input.""" + annotator = HeatMapAnnotator() + detections = _create_detections(xyxy=[[20, 20, 60, 60]]) + result = annotator.annotate(scene=test_image.copy(), detections=detections) + assert not np.array_equal(test_image, result) + + def test_annotate_state_preserved_after_empty_call( + self, test_image: np.ndarray + ) -> None: + """Empty call must not poison accumulated heat.""" + annotator = HeatMapAnnotator() + detections = _create_detections(xyxy=[[20, 20, 60, 60]]) + annotator.annotate(scene=test_image.copy(), detections=Detections.empty()) + result = annotator.annotate(scene=test_image.copy(), detections=detections) + assert not np.array_equal(test_image, result) + + def test_annotate_empty_after_real_does_not_warn( + self, test_image: np.ndarray + ) -> None: + """Empty call after heat accumulated must not trigger RuntimeWarning.""" + annotator = HeatMapAnnotator() + detections = _create_detections(xyxy=[[20, 20, 60, 60]]) + annotator.annotate(scene=test_image.copy(), detections=detections) + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + annotator.annotate(scene=test_image.copy(), detections=Detections.empty()) + + class TestEllipseAnnotator: """Tests for EllipseAnnotator class"""