perf(annotators): share mask painting and give HaloAnnotator the CompactMask path (#2339)
MaskAnnotator paints CompactMask detections into their bounding-box crop, but HaloAnnotator never got that path: it materialized every mask full-frame, painted via full-frame boolean indexing, and built its foreground mask as a 2M-element Python list per frame. Extract the shared compact/dense painting into a single _paint_masks_by_area helper used by both annotators. On a 1080p frame with 30 masks, HaloAnnotator on CompactMask runs about 4x faster; output is unchanged. - Replace in-place `union` param with `collect_union: bool` return value; HaloAnnotator now captures the returned union array - Add Google-style Args/Returns to `_paint_masks_by_area` docstring - Fix HaloAnnotator.annotate() example (was maskless → no-op) and scene arg wording - Add section comment above shared helper for discoverability - Add union accumulation tests (dense + CompactMask paths via collect_union=True) - Add test documenting out-of-bbox True-pixel divergence between compact and dense - Add image-edge bbox test for CompactMask annotators - Move helper tests into TestPaintMasksByArea and TestCompactMaskParity classes - Rename test_annotate_with_empty_masks → test_annotate_with_all_false_mask --------- Co-authored-by: jirka <6035284+Borda@users.noreply.github.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
This commit is contained in:
Agis Kounelis
GitHub
parent
11a586c133
commit
a48532f6f5
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
Loading…
Reference in New Issue