diff --git a/src/supervision/annotators/core.py b/src/supervision/annotators/core.py index 559528cc..8336e11a 100644 --- a/src/supervision/annotators/core.py +++ b/src/supervision/annotators/core.py @@ -361,6 +361,65 @@ class OrientedBoxAnnotator(BaseAnnotator): return scene +# --- Shared mask-painting utilities --- +def _paint_masks_by_area( + canvas: npt.NDArray[np.uint8], + detections: Detections, + color: Color | ColorPalette, + color_lookup: ColorLookup | npt.NDArray[np.int_], + collect_union: bool = False, +) -> npt.NDArray[np.bool_] | None: + """Paint each detection's mask into `canvas` in descending-area order. + + Smaller masks are drawn on top of larger ones. `CompactMask` detections + are painted into their bounding-box crop only, avoiding a full `(H, W)` + allocation per mask; dense masks fall back to full-frame boolean indexing. + + Args: + canvas: BGR image array painted in place. Shape ``(H, W, 3)``. + detections: Detections whose masks to paint. Returns immediately + without modifying `canvas` when ``detections.mask`` is ``None``. + color: Single color or palette used to resolve each detection's color. + color_lookup: Strategy for mapping colors to detection indices. + collect_union: When ``True``, allocate and return a ``(H, W)`` + boolean array that accumulates the union of all painted masks + (useful for callers like `HaloAnnotator` that need the combined + mask footprint). When ``False`` (default), returns ``None``. + + Returns: + A ``(H, W)`` boolean union array when ``collect_union=True``, + otherwise ``None``. + """ + masks = detections.mask + if masks is None: + return None + union: npt.NDArray[np.bool_] | None = ( + np.zeros(canvas.shape[:2], dtype=bool) if collect_union else None + ) + compact_mask = masks if isinstance(masks, CompactMask) else None + for detection_idx in np.flip(np.argsort(detections.area)): + color_bgr = resolve_color( + color=color, + detections=detections, + detection_idx=detection_idx, + color_lookup=color_lookup, + ).as_bgr() + if compact_mask is not None: + x1 = int(compact_mask.offsets[detection_idx, 0]) + y1 = int(compact_mask.offsets[detection_idx, 1]) + crop_m = compact_mask.crop(detection_idx) + crop_h, crop_w = crop_m.shape + canvas[y1 : y1 + crop_h, x1 : x1 + crop_w][crop_m] = color_bgr + if union is not None: + union[y1 : y1 + crop_h, x1 : x1 + crop_w] |= crop_m + else: + mask = np.asarray(masks[detection_idx], dtype=bool) + canvas[mask] = color_bgr + if union is not None: + union |= mask + return union + + class MaskAnnotator(BaseAnnotator): """ A class for drawing masks on an image using provided detections. @@ -437,35 +496,12 @@ class MaskAnnotator(BaseAnnotator): return scene colored_mask = np.array(scene, copy=True, dtype=np.uint8) - - compact_mask = ( - detections.mask if isinstance(detections.mask, CompactMask) else None + _paint_masks_by_area( + colored_mask, + detections, + self.color, + self.color_lookup if custom_color_lookup is None else custom_color_lookup, ) - for detection_idx in np.flip(np.argsort(detections.area)): - color = resolve_color( - color=self.color, - detections=detections, - detection_idx=detection_idx, - color_lookup=self.color_lookup - if custom_color_lookup is None - else custom_color_lookup, - ) - if compact_mask is not None: - # Paint only the bounding-box crop — avoids a full (H, W) alloc. - x1 = int(compact_mask.offsets[detection_idx, 0]) - y1 = int(compact_mask.offsets[detection_idx, 1]) - crop_m = compact_mask.crop(detection_idx) - crop_h, crop_w = crop_m.shape - colored_mask[y1 : y1 + crop_h, x1 : x1 + crop_w][crop_m] = ( - color.as_bgr() - ) - else: - mask = np.asarray( - detections.mask[detection_idx], - dtype=bool, - ) - colored_mask[mask] = color.as_bgr() - cv2.addWeighted( colored_mask, self.opacity, scene, 1 - self.opacity, 0, dst=scene ) @@ -701,7 +737,7 @@ class HaloAnnotator(BaseAnnotator): Annotates the given scene with halos based on the provided detections. Args: - scene: The image where masks will be drawn. + scene: The image where the halo effect will be applied. `ImageType` is a flexible type, accepting either `numpy.ndarray` or `PIL.Image.Image`. detections: Object detections to annotate. @@ -719,6 +755,7 @@ class HaloAnnotator(BaseAnnotator): >>> image = np.zeros((100, 100, 3), dtype=np.uint8) >>> detections = sv.Detections( ... xyxy=np.array([[20, 20, 80, 80]]), + ... mask=np.zeros((1, 100, 100), dtype=bool), ... class_id=np.array([0]) ... ) >>> halo_annotator = sv.HaloAnnotator() @@ -737,28 +774,23 @@ class HaloAnnotator(BaseAnnotator): if detections.mask is None: return scene colored_mask = np.zeros_like(scene, dtype=np.uint8) - fmask = np.array([False] * scene.shape[0] * scene.shape[1]).reshape( - scene.shape[0], scene.shape[1] + fmask = _paint_masks_by_area( + colored_mask, + detections, + self.color, + self.color_lookup if custom_color_lookup is None else custom_color_lookup, + collect_union=True, ) - - for detection_idx in np.flip(np.argsort(detections.area)): - color = resolve_color( - color=self.color, - detections=detections, - detection_idx=detection_idx, - color_lookup=self.color_lookup - if custom_color_lookup is None - else custom_color_lookup, - ) - mask = np.asarray(detections.mask[detection_idx], dtype=bool) - fmask = np.logical_or(fmask, mask) - color_bgr = color.as_bgr() - colored_mask[mask] = color_bgr + assert fmask is not None # collect_union=True always returns an array colored_mask = cv2.blur(colored_mask, (self.kernel_size, self.kernel_size)) colored_mask[fmask] = [0, 0, 0] gray = cv2.cvtColor(colored_mask, cv2.COLOR_BGR2GRAY) - alpha = self.opacity * gray / gray.max() + gray_max = gray.max() + if gray_max == 0: + # no halo to draw (e.g. empty masks); leave the scene untouched + return scene + alpha = self.opacity * gray / gray_max alpha_mask = alpha[:, :, np.newaxis] blended_scene = np.uint8(scene * (1 - alpha_mask) + colored_mask * self.opacity) np.copyto(scene, blended_scene) diff --git a/tests/annotators/test_core.py b/tests/annotators/test_core.py index 51428edb..005f8166 100644 --- a/tests/annotators/test_core.py +++ b/tests/annotators/test_core.py @@ -30,8 +30,10 @@ from supervision.annotators.core import ( RoundBoxAnnotator, TraceAnnotator, TriangleAnnotator, + _paint_masks_by_area, ) from supervision.annotators.utils import ColorLookup +from supervision.detection.compact_mask import CompactMask from supervision.detection.core import Detections from supervision.draw.color import Color from supervision.geometry.core import Position @@ -363,6 +365,173 @@ class TestHaloAnnotator: ) assert np.array_equal(result_bool, result_uint8) + def test_annotate_with_all_false_mask_preserves_scene(self): + """Test that an all-False mask leaves the scene unchanged, not corrupted.""" + scene = np.full((100, 100, 3), 127, dtype=np.uint8) + masks = [np.zeros((100, 100), dtype=bool)] + detections = _create_detections( + xyxy=[[10, 10, 90, 90]], mask=masks, class_id=[0] + ) + result = HaloAnnotator().annotate(scene=scene.copy(), detections=detections) + assert np.array_equal(result, scene) + + +class TestPaintMasksByArea: + """Tests for the _paint_masks_by_area helper function.""" + + def test_paint_masks_by_area_is_noop_without_masks(self): + """_paint_masks_by_area is a no-op when detections carry no mask.""" + canvas = np.full((10, 10, 3), 7, dtype=np.uint8) + detections = _create_detections(xyxy=[[1, 1, 8, 8]], class_id=[0]) + _paint_masks_by_area(canvas, detections, Color.RED, ColorLookup.INDEX) + assert np.array_equal(canvas, np.full((10, 10, 3), 7, dtype=np.uint8)) + + def test_union_accumulation_dense(self): + """Dense path: collect_union=True returns array covering all painted pixels.""" + height, width = 50, 60 + canvas = np.zeros((height, width, 3), dtype=np.uint8) + masks = [np.zeros((height, width), dtype=bool)] + masks[0][5:20, 10:40] = True + detections = _create_detections( + xyxy=[[10.0, 5.0, 40.0, 20.0]], mask=masks, class_id=[0] + ) + result_union = _paint_masks_by_area( + canvas, detections, Color.RED, ColorLookup.INDEX, collect_union=True + ) + assert result_union is not None + # every painted pixel must be in the union (RED is BGR (0, 0, 255), + # so detect painted pixels via any non-zero channel) + painted = canvas.any(axis=-1) + assert np.array_equal(painted, result_union) + + def test_union_accumulation_compact(self): + """CompactMask path: collect_union=True returns array matching dense.""" + height, width = 50, 60 + mask = np.zeros((height, width), dtype=bool) + mask[5:20, 10:40] = True + xyxy = np.array([[10.0, 5.0, 40.0, 20.0]]) + + canvas_dense = np.zeros((height, width, 3), dtype=np.uint8) + dense = _create_detections(xyxy=xyxy.tolist(), mask=[mask], class_id=[0]) + union_dense = _paint_masks_by_area( + canvas_dense, dense, Color.RED, ColorLookup.INDEX, collect_union=True + ) + + canvas_compact = np.zeros((height, width, 3), dtype=np.uint8) + compact = _create_detections(xyxy=xyxy.tolist(), mask=[mask], class_id=[0]) + compact.mask = CompactMask.from_dense( + np.array([mask]), compact.xyxy, (height, width) + ) + union_compact = _paint_masks_by_area( + canvas_compact, compact, Color.RED, ColorLookup.INDEX, collect_union=True + ) + + assert union_dense is not None + assert union_compact is not None + # compact union must cover exactly the same pixels as dense union + assert np.array_equal(union_dense, union_compact) + + def test_compact_mask_drops_pixels_outside_bbox(self): + """CompactMask is lossy: True pixels outside xyxy bbox are silently dropped. + + This test documents that compact and dense paths diverge when a mask has + True pixels outside its bounding box — the 'bit-identical' claim holds + only for bbox-contained masks. + """ + height, width = 50, 60 + mask = np.zeros((height, width), dtype=bool) + mask[5:25, 10:40] = True # mask extends 5 rows beyond bbox bottom + + bbox = [[10.0, 5.0, 40.0, 20.0]] # y2=20 clips the mask at row 20 + + canvas_dense = np.zeros((height, width, 3), dtype=np.uint8) + dense = _create_detections(xyxy=bbox, mask=[mask], class_id=[0]) + _paint_masks_by_area(canvas_dense, dense, Color.RED, ColorLookup.INDEX) + + canvas_compact = np.zeros((height, width, 3), dtype=np.uint8) + compact = _create_detections(xyxy=bbox, mask=[mask], class_id=[0]) + compact.mask = CompactMask.from_dense( + np.array([mask]), compact.xyxy, (height, width) + ) + _paint_masks_by_area(canvas_compact, compact, Color.RED, ColorLookup.INDEX) + + # Dense paints all True pixels incl. rows 21-24; compact only within bbox. + assert not np.array_equal(canvas_dense, canvas_compact), ( + "Expected divergence: compact mask drops True pixels outside bbox" + ) + # Compact subset: every pixel painted by compact is also painted by dense. + compact_painted = canvas_compact.any(axis=-1) + dense_painted = canvas_dense.any(axis=-1) + assert np.all(dense_painted[compact_painted]) + + +class TestCompactMaskParity: + """Tests that CompactMask and dense mask produce identical annotator output.""" + + @pytest.mark.parametrize( + "annotator_factory", + [ + pytest.param( + lambda: MaskAnnotator(opacity=1.0, color_lookup=ColorLookup.INDEX), + id="mask", + ), + pytest.param( + lambda: HaloAnnotator(kernel_size=15, color_lookup=ColorLookup.INDEX), + id="halo", + ), + ], + ) + def test_annotator_compact_mask_matches_dense_mask(self, annotator_factory): + """CompactMask detections annotate identically to dense bool masks.""" + height, width = 120, 160 + rng = np.random.default_rng(0) + scene = rng.integers(0, 256, (height, width, 3), dtype=np.uint8) + boxes = [[10, 10, 70, 60], [40, 30, 150, 110], [90, 70, 140, 115]] + masks = [] + for x1, y1, x2, y2 in boxes: + mask = np.zeros((height, width), dtype=bool) + mask[y1 : y2 + 1, x1 : x2 + 1] = True + masks.append(mask) + class_id = [0, 1, 2] + xyxy = [[float(value) for value in box] for box in boxes] + + dense = _create_detections(xyxy=xyxy, mask=masks, class_id=class_id) + compact = _create_detections(xyxy=xyxy, mask=masks, class_id=class_id) + compact.mask = CompactMask.from_dense( + np.array(masks), compact.xyxy, (height, width) + ) + + result_dense = annotator_factory().annotate( + scene=scene.copy(), detections=dense + ) + result_compact = annotator_factory().annotate( + scene=scene.copy(), detections=compact + ) + + assert not np.array_equal(result_dense, scene), "annotator painted nothing" + assert np.array_equal(result_dense, result_compact) + + def test_annotator_compact_mask_handles_edge_clipping(self): + """CompactMask detection straddling image edge paints via NumPy clip.""" + height, width = 50, 60 + rng = np.random.default_rng(42) + scene = rng.integers(0, 256, (height, width, 3), dtype=np.uint8) + + # Box extends 10 pixels beyond right/bottom edges + mask = np.zeros((height, width), dtype=bool) + mask[40:height, 50:width] = True + bbox = [[50.0, 40.0, width + 10.0, height + 10.0]] + + detections = _create_detections(xyxy=bbox, mask=[mask], class_id=[0]) + detections.mask = CompactMask.from_dense( + np.array([mask]), detections.xyxy, (height, width) + ) + + annotator = MaskAnnotator(opacity=1.0, color_lookup=ColorLookup.INDEX) + result = annotator.annotate(scene=scene.copy(), detections=detections) + # Result must differ from scene (something was painted) and must not raise + assert not np.array_equal(result, scene), "Expected pixels to be painted" + class TestHeatMapAnnotator: """Tests for HeatMapAnnotator class"""