wehub-resource-sync
9194ef5abd
Docs/Test Workflow / Test docs build (push) Failing after 0s
Check links & references / links-check (push) Failing after 1s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.10) (push) Failing after 0s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.11) (push) Failing after 0s
PR Conflict Labeler / main (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.12) (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.13) (push) Failing after 0s
Pytest/Test Workflow / Build this Package (push) Failing after 5s
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / testing-guardian (push) Has been cancelled
1960 lines
84 KiB
Python
1960 lines
84 KiB
Python
"""
|
|
Tests for supervision/annotators/core.py
|
|
"""
|
|
|
|
import warnings
|
|
from collections.abc import Iterator
|
|
from typing import Any, cast
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
import supervision.annotators.core as annotators_core
|
|
from supervision.annotators.base import BaseAnnotator
|
|
from supervision.annotators.core import (
|
|
BackgroundOverlayAnnotator,
|
|
BlurAnnotator,
|
|
BoxAnnotator,
|
|
BoxCornerAnnotator,
|
|
CircleAnnotator,
|
|
ColorAnnotator,
|
|
ComparisonAnnotator,
|
|
CropAnnotator,
|
|
DotAnnotator,
|
|
EllipseAnnotator,
|
|
HaloAnnotator,
|
|
HeatMapAnnotator,
|
|
IconAnnotator,
|
|
LabelAnnotator,
|
|
MaskAnnotator,
|
|
OrientedBoxAnnotator,
|
|
PercentageBarAnnotator,
|
|
PixelateAnnotator,
|
|
PolygonAnnotator,
|
|
RichLabelAnnotator,
|
|
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
|
|
from tests.helpers import _create_detections, assert_image_mostly_same
|
|
|
|
|
|
def _get_concrete_annotator_subclasses(
|
|
cls: type[BaseAnnotator],
|
|
) -> Iterator[type[BaseAnnotator]]:
|
|
"""Recursively yield non-abstract BaseAnnotator subclasses."""
|
|
for sub in cls.__subclasses__():
|
|
if not getattr(sub, "__abstractmethods__", None):
|
|
yield sub
|
|
yield from _get_concrete_annotator_subclasses(sub)
|
|
|
|
|
|
class TestAnnotatorMaskPolicy:
|
|
"""Tests for annotator mask materialization policy metadata."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"annotator",
|
|
[
|
|
pytest.param(MaskAnnotator(), id="mask"),
|
|
pytest.param(PolygonAnnotator(), id="polygon"),
|
|
pytest.param(HaloAnnotator(), id="halo"),
|
|
],
|
|
)
|
|
def test_mask_required_annotators_declare_mask_requirement(self, annotator):
|
|
"""Annotators that require detections.mask expose the requirement."""
|
|
assert type(annotator).requires_mask is True
|
|
|
|
@pytest.mark.parametrize(
|
|
"annotator",
|
|
[
|
|
pytest.param(BoxAnnotator(), id="box"),
|
|
pytest.param(LabelAnnotator(), id="label"),
|
|
pytest.param(CircleAnnotator(), id="circle"),
|
|
pytest.param(EllipseAnnotator(), id="ellipse"),
|
|
pytest.param(IconAnnotator(), id="icon"),
|
|
pytest.param(TraceAnnotator(), id="trace"),
|
|
pytest.param(BackgroundOverlayAnnotator(), id="overlay"),
|
|
pytest.param(BackgroundOverlayAnnotator(force_box=True), id="overlay-box"),
|
|
pytest.param(ComparisonAnnotator(), id="comparison"),
|
|
],
|
|
)
|
|
def test_mask_optional_annotators_declare_no_mask_requirement(self, annotator):
|
|
"""Mask-optional annotators do not require mask materialization."""
|
|
assert type(annotator).requires_mask is False
|
|
|
|
@pytest.mark.parametrize(
|
|
"annotator_class",
|
|
[
|
|
pytest.param(cls, id=cls.__name__)
|
|
for cls in _get_concrete_annotator_subclasses(BaseAnnotator)
|
|
],
|
|
)
|
|
def test_all_subclasses_have_bool_requires_mask(self, annotator_class):
|
|
"""Every concrete BaseAnnotator subclass declares requires_mask as a bool."""
|
|
assert isinstance(annotator_class.requires_mask, bool)
|
|
|
|
def test_exact_mask_requiring_annotator_set(self):
|
|
"""Only MaskAnnotator, PolygonAnnotator, HaloAnnotator require masks."""
|
|
mask_true = {
|
|
cls
|
|
for cls in _get_concrete_annotator_subclasses(BaseAnnotator)
|
|
if cls.requires_mask is True
|
|
}
|
|
assert mask_true == {MaskAnnotator, PolygonAnnotator, HaloAnnotator}
|
|
|
|
|
|
@pytest.fixture
|
|
def test_image() -> np.ndarray:
|
|
"""Create a simple blank test image fixture"""
|
|
return np.zeros((100, 100, 3), dtype=np.uint8)
|
|
|
|
|
|
@pytest.fixture
|
|
def test_mask() -> np.ndarray:
|
|
"""Create a simple rectangular mask fixture"""
|
|
mask = np.zeros((100, 100), dtype=bool)
|
|
mask[20:80, 20:80] = True
|
|
return mask
|
|
|
|
|
|
@pytest.fixture
|
|
def gradient_image() -> np.ndarray:
|
|
"""Create a gradient test image fixture"""
|
|
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
|
for i in range(100):
|
|
for j in range(100):
|
|
image[i, j] = [i, j, (i + j) // 2]
|
|
return image
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("factory", "expected_colors"),
|
|
[
|
|
(lambda: BoxAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(lambda: OrientedBoxAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(lambda: MaskAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(lambda: PolygonAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(lambda: ColorAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(lambda: HaloAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(lambda: EllipseAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(lambda: BoxCornerAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(lambda: CircleAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(
|
|
lambda: DotAnnotator(color="#010203", outline_color="#040506"),
|
|
{"color": (1, 2, 3), "outline_color": (4, 5, 6)},
|
|
),
|
|
(
|
|
lambda: LabelAnnotator(color="#010203", text_color="#040506"),
|
|
{"color": (1, 2, 3), "text_color": (4, 5, 6)},
|
|
),
|
|
(
|
|
lambda: RichLabelAnnotator(color="#010203", text_color="#040506"),
|
|
{"color": (1, 2, 3), "text_color": (4, 5, 6)},
|
|
),
|
|
(lambda: TraceAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(
|
|
lambda: TriangleAnnotator(color="#010203", outline_color="#040506"),
|
|
{"color": (1, 2, 3), "outline_color": (4, 5, 6)},
|
|
),
|
|
(lambda: RoundBoxAnnotator(color="#010203"), {"color": (1, 2, 3)}),
|
|
(
|
|
lambda: PercentageBarAnnotator(color="#010203", border_color="#040506"),
|
|
{"color": (1, 2, 3), "border_color": (4, 5, 6)},
|
|
),
|
|
(lambda: CropAnnotator(border_color="#010203"), {"border_color": (1, 2, 3)}),
|
|
],
|
|
)
|
|
def test_hex_color_support_across_annotators(
|
|
factory, expected_colors: dict[str, tuple[int, int, int]]
|
|
) -> None:
|
|
annotator = factory()
|
|
for attribute_name, expected_rgb in expected_colors.items():
|
|
color = getattr(annotator, attribute_name)
|
|
assert isinstance(color, Color)
|
|
assert color.as_rgb() == expected_rgb
|
|
|
|
|
|
class TestBoxAnnotator:
|
|
"""
|
|
Verify that BoxAnnotator correctly draws bounding boxes on an image.
|
|
|
|
Ensures that `BoxAnnotator` correctly draws bounding boxes on an image, which is
|
|
essential for users to visualize detection results.
|
|
"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image: np.ndarray) -> None:
|
|
"""
|
|
Verify that annotation with no detections does not change the image.
|
|
|
|
Scenario: Annotating an image with an empty set of detections.
|
|
Expected: The scene remains unchanged, ensuring no ghost boxes are drawn.
|
|
"""
|
|
detections = Detections.empty()
|
|
annotator = BoxAnnotator()
|
|
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:
|
|
"""
|
|
Verify that annotation with a single detection draws a bounding box.
|
|
|
|
Scenario: Annotating an image with a single bounding box.
|
|
Expected: The scene is modified by drawing a box, allowing users to identify
|
|
a single detected object.
|
|
"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = BoxAnnotator(
|
|
color=Color.WHITE, thickness=2, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.85)
|
|
|
|
def test_annotate_with_multiple_detections(self, test_image: np.ndarray) -> None:
|
|
"""
|
|
Verify that annotation with multiple detections draws all bounding boxes.
|
|
|
|
Scenario: Annotating an image with multiple bounding boxes of different classes.
|
|
Expected: All boxes are drawn, enabling visualization of complex scenes with
|
|
multiple objects.
|
|
"""
|
|
detections = _create_detections(
|
|
xyxy=[[10, 10, 40, 40], [60, 60, 90, 90], [10, 60, 40, 90]],
|
|
class_id=[0, 1, 2],
|
|
)
|
|
annotator = BoxAnnotator(
|
|
color=Color.WHITE, thickness=2, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.85)
|
|
|
|
def test_annotate_with_numpy_color_lookup(self, test_image: np.ndarray) -> None:
|
|
"""
|
|
Verify that annotation respects custom NumPy color lookup array.
|
|
|
|
Scenario: Providing a custom NumPy array for color lookup instead of class IDs.
|
|
Expected: Annotator respects the custom mapping, giving users flexible control
|
|
over box colors (e.g., coloring by tracking ID or custom criteria).
|
|
"""
|
|
detections = Detections(
|
|
xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40]], dtype=np.float32),
|
|
confidence=np.array([0.38, 0.21], dtype=np.float32),
|
|
class_id=np.array([0, 0], dtype=np.int64),
|
|
tracker_id=None,
|
|
)
|
|
|
|
lookup = np.array([1, 0], dtype=np.int16)
|
|
|
|
annotator = BoxAnnotator(
|
|
color=Color.WHITE, thickness=2, color_lookup=ColorLookup.INDEX
|
|
)
|
|
|
|
result = annotator.annotate(
|
|
scene=test_image.copy(),
|
|
detections=detections,
|
|
custom_color_lookup=lookup,
|
|
)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.85)
|
|
|
|
|
|
class TestOrientedBoxAnnotator:
|
|
"""Tests for OrientedBoxAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = OrientedBoxAnnotator()
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_without_oriented_boxes(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no OBB data"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]])
|
|
annotator = OrientedBoxAnnotator()
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
|
|
class TestMaskAnnotator:
|
|
"""Tests for MaskAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = MaskAnnotator()
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_without_masks(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no masks"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = MaskAnnotator(color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_with_single_mask(self, test_image, test_mask):
|
|
"""Test that annotate method correctly draws a single mask"""
|
|
detections = _create_detections(
|
|
xyxy=[[10, 10, 90, 90]], mask=[test_mask], class_id=[0]
|
|
)
|
|
annotator = MaskAnnotator(
|
|
color=Color.RED, opacity=1.0, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.6)
|
|
|
|
def test_annotate_uint8_mask_matches_bool_mask(self, test_image, test_mask):
|
|
"""Test that uint8 and bool masks produce identical overlays."""
|
|
detections_bool = _create_detections(
|
|
xyxy=[[10, 10, 90, 90]], mask=[test_mask], class_id=[0]
|
|
)
|
|
detections_uint8 = _create_detections(
|
|
xyxy=[[10, 10, 90, 90]], mask=[test_mask], class_id=[0]
|
|
)
|
|
detections_uint8.mask = detections_uint8.mask.astype(np.uint8)
|
|
|
|
annotator = MaskAnnotator(
|
|
color=Color.RED, opacity=1.0, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result_bool = annotator.annotate(
|
|
scene=test_image.copy(), detections=detections_bool
|
|
)
|
|
result_uint8 = annotator.annotate(
|
|
scene=test_image.copy(), detections=detections_uint8
|
|
)
|
|
assert np.array_equal(result_bool, result_uint8)
|
|
|
|
def test_annotate_blends_only_dense_mask_roi(self, monkeypatch):
|
|
"""Dense mask annotation should blend only the touched ROI."""
|
|
height, width = 80, 90
|
|
scene = np.random.default_rng(1).integers(
|
|
0, 256, (height, width, 3), dtype=np.uint8
|
|
)
|
|
masks = np.zeros((2, height, width), dtype=bool)
|
|
masks[0, 10:20, 15:25] = True
|
|
masks[1, 45:55, 60:75] = True
|
|
# inclusive xyxy; floor(x2)+1 gives union ROI (15,10,75,55) → shape (45,60,3)
|
|
detections = _create_detections(
|
|
xyxy=[[15.0, 10.0, 24.0, 19.0], [60.0, 45.0, 74.0, 54.0]],
|
|
mask=masks,
|
|
class_id=[0, 1],
|
|
)
|
|
original_add_weighted = cv2.addWeighted
|
|
blended_shapes = []
|
|
|
|
def add_weighted_spy(src1, alpha, src2, beta, gamma, dst=None, dtype=None):
|
|
blended_shapes.append(src1.shape)
|
|
return original_add_weighted(src1, alpha, src2, beta, gamma, dst, dtype)
|
|
|
|
monkeypatch.setattr(cv2, "addWeighted", add_weighted_spy)
|
|
|
|
result = MaskAnnotator(
|
|
opacity=0.5, color=Color.RED, color_lookup=ColorLookup.INDEX
|
|
).annotate(scene=scene.copy(), detections=detections)
|
|
|
|
assert not np.array_equal(result, scene)
|
|
assert blended_shapes == [(45, 60, 3)]
|
|
|
|
def test_annotate_blends_only_compact_mask_roi(self, monkeypatch):
|
|
"""CompactMask annotation should blend only the touched ROI."""
|
|
height, width = 80, 90
|
|
scene = np.random.default_rng(2).integers(
|
|
0, 256, (height, width, 3), dtype=np.uint8
|
|
)
|
|
masks = np.zeros((2, height, width), dtype=bool)
|
|
masks[0, 10:20, 15:25] = True
|
|
masks[1, 45:55, 60:75] = True
|
|
xyxy = np.array([[15.0, 10.0, 24.0, 19.0], [60.0, 45.0, 74.0, 54.0]])
|
|
detections = _create_detections(xyxy=xyxy.tolist(), mask=masks, class_id=[0, 1])
|
|
detections.mask = CompactMask.from_dense(
|
|
masks, detections.xyxy, (height, width)
|
|
)
|
|
original_add_weighted = cv2.addWeighted
|
|
blended_shapes = []
|
|
|
|
def add_weighted_spy(src1, alpha, src2, beta, gamma, dst=None, dtype=None):
|
|
blended_shapes.append(src1.shape)
|
|
return original_add_weighted(src1, alpha, src2, beta, gamma, dst, dtype)
|
|
|
|
monkeypatch.setattr(cv2, "addWeighted", add_weighted_spy)
|
|
|
|
result = MaskAnnotator(opacity=0.5, color_lookup=ColorLookup.INDEX).annotate(
|
|
scene=scene.copy(), detections=detections
|
|
)
|
|
|
|
assert not np.array_equal(result, scene)
|
|
assert blended_shapes == [(45, 60, 3)]
|
|
|
|
def test_annotate_skips_all_false_mask_blend(self, monkeypatch):
|
|
"""All-false masks must skip blending — addWeighted must not be called."""
|
|
height, width = 30, 40
|
|
scene = np.random.default_rng(3).integers(
|
|
0, 256, (height, width, 3), dtype=np.uint8
|
|
)
|
|
mask = np.zeros((height, width), dtype=bool)
|
|
detections = _create_detections(
|
|
xyxy=[[5.0, 5.0, 20.0, 20.0]], mask=[mask], class_id=[0]
|
|
)
|
|
|
|
call_count = []
|
|
|
|
def add_weighted_spy(src1, alpha, src2, beta, gamma, dst=None, dtype=None):
|
|
call_count.append(1)
|
|
return src2
|
|
|
|
monkeypatch.setattr(cv2, "addWeighted", add_weighted_spy)
|
|
|
|
result = MaskAnnotator(
|
|
opacity=0.5, color=Color.RED, color_lookup=ColorLookup.INDEX
|
|
).annotate(scene=scene.copy(), detections=detections)
|
|
|
|
assert np.array_equal(result, scene)
|
|
assert len(call_count) == 0, "addWeighted called for all-false masks"
|
|
|
|
def test_annotate_pixels_outside_roi_unchanged(self):
|
|
"""Pixels outside the blended ROI must be unchanged from the original scene."""
|
|
height, width = 80, 90
|
|
rng = np.random.default_rng(42)
|
|
scene = rng.integers(0, 256, (height, width, 3), dtype=np.uint8)
|
|
# Two masks in the top-left region — ROI should be [10:55, 15:75]
|
|
masks = np.zeros((2, height, width), dtype=bool)
|
|
masks[0, 10:20, 15:25] = True
|
|
masks[1, 45:55, 60:75] = True
|
|
# inclusive xyxy; floor(x2)+1 gives ROI (15,10,75,55)
|
|
xyxy = np.array([[15.0, 10.0, 24.0, 19.0], [60.0, 45.0, 74.0, 54.0]])
|
|
detections = _create_detections(xyxy=xyxy.tolist(), mask=masks, class_id=[0, 1])
|
|
result = MaskAnnotator(opacity=0.5, color_lookup=ColorLookup.INDEX).annotate(
|
|
scene=scene.copy(), detections=detections
|
|
)
|
|
# Pixels strictly below ROI row-bound (y2=55) must be unchanged
|
|
assert np.array_equal(result[55:, :], scene[55:, :])
|
|
# Pixels strictly right of ROI col-bound (x2=75) must be unchanged
|
|
assert np.array_equal(result[:, 75:], scene[:, 75:])
|
|
|
|
def test_annotate_roi_parity_with_full_frame_blend(self):
|
|
"""ROI blend must match reference blend within ±1 (uint8 rounding)."""
|
|
height, width = 60, 80
|
|
rng = np.random.default_rng(99)
|
|
scene = rng.integers(0, 256, (height, width, 3), dtype=np.uint8)
|
|
opacity = 0.5
|
|
# Single mask in bottom-right quadrant
|
|
masks = np.zeros((1, height, width), dtype=bool)
|
|
masks[0, 40:55, 50:70] = True
|
|
# inclusive xyxy; floor(x2)+1 gives ROI scene[40:55, 50:70]
|
|
xyxy = np.array([[50.0, 40.0, 69.0, 54.0]])
|
|
detections = _create_detections(xyxy=xyxy.tolist(), mask=masks, class_id=[0])
|
|
|
|
# ROI-only result (new optimized behavior)
|
|
result_roi = MaskAnnotator(
|
|
opacity=opacity, color=Color.RED, color_lookup=ColorLookup.INDEX
|
|
).annotate(scene=scene.copy(), detections=detections)
|
|
|
|
# Reference: manually compute expected blend for masked pixels only
|
|
# RED in BGR = (0, 0, 255); blend = opacity*color + (1-opacity)*scene
|
|
roi_mask = masks[0]
|
|
colored = scene.copy()
|
|
colored[roi_mask] = (0, 0, 255) # BGR RED
|
|
ref_blend = np.clip(
|
|
opacity * colored.astype(np.float32)
|
|
+ (1 - opacity) * scene.astype(np.float32),
|
|
0,
|
|
255,
|
|
).astype(np.uint8)
|
|
|
|
# Masked pixels must match reference within ±1 (uint8 rounding)
|
|
diff = np.abs(
|
|
result_roi[roi_mask].astype(np.int16) - ref_blend[roi_mask].astype(np.int16)
|
|
)
|
|
assert np.all(diff <= 1), f"Max pixel diff: {diff.max()}"
|
|
|
|
|
|
class TestPolygonAnnotator:
|
|
"""Tests for PolygonAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = PolygonAnnotator()
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_without_masks(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no masks"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = PolygonAnnotator(color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_with_single_mask(self, test_image, test_mask):
|
|
"""Test that annotate method correctly draws a single polygon from mask"""
|
|
detections = _create_detections(
|
|
xyxy=[[10, 10, 90, 90]], mask=[test_mask], class_id=[0]
|
|
)
|
|
annotator = PolygonAnnotator(
|
|
color=Color.WHITE, thickness=2, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.85)
|
|
|
|
def test_compact_mask_uses_crops_and_matches_dense_mask(self, monkeypatch):
|
|
"""CompactMask polygons use crops without full-frame mask indexing."""
|
|
height, width = 80, 90
|
|
scene = np.zeros((height, width, 3), dtype=np.uint8)
|
|
masks = np.zeros((3, height, width), dtype=bool)
|
|
masks[0, 10:20, 15:30] = True
|
|
masks[0, 32:45, 40:55] = True
|
|
masks[1] = False
|
|
masks[2, 50:70, 5:25] = True
|
|
xyxy = np.array(
|
|
[[15, 10, 54, 44], [60, 5, 70, 15], [5, 50, 24, 69]],
|
|
dtype=np.float32,
|
|
)
|
|
dense = Detections(xyxy=xyxy, mask=masks, class_id=np.array([0, 1, 2]))
|
|
compact = Detections(
|
|
xyxy=xyxy,
|
|
mask=CompactMask.from_dense(masks, xyxy, (height, width)),
|
|
class_id=np.array([0, 1, 2]),
|
|
)
|
|
annotator = PolygonAnnotator(
|
|
color=Color.WHITE, thickness=1, color_lookup=ColorLookup.INDEX
|
|
)
|
|
expected = annotator.annotate(scene=scene.copy(), detections=dense)
|
|
|
|
original_getitem = CompactMask.__getitem__
|
|
|
|
def fail_int_index(self, index):
|
|
if isinstance(index, (int, np.integer)):
|
|
raise AssertionError("PolygonAnnotator must use CompactMask.crop")
|
|
return original_getitem(self, index)
|
|
|
|
def fail_to_dense(self):
|
|
raise AssertionError("PolygonAnnotator must not call CompactMask.to_dense")
|
|
|
|
monkeypatch.setattr(CompactMask, "__getitem__", fail_int_index)
|
|
monkeypatch.setattr(CompactMask, "to_dense", fail_to_dense)
|
|
result = annotator.annotate(scene=scene.copy(), detections=compact)
|
|
|
|
assert not np.array_equal(result, scene), "annotator painted nothing"
|
|
np.testing.assert_array_equal(result, expected)
|
|
|
|
def test_annotate_with_empty_compact_mask(self):
|
|
"""N=0 CompactMask returns scene unchanged without error."""
|
|
height, width = 60, 80
|
|
scene = np.zeros((height, width, 3), dtype=np.uint8)
|
|
empty_masks = np.zeros((0, height, width), dtype=bool)
|
|
xyxy = np.zeros((0, 4), dtype=np.float32)
|
|
detections = Detections(
|
|
xyxy=xyxy,
|
|
mask=CompactMask.from_dense(empty_masks, xyxy, (height, width)),
|
|
class_id=np.array([], dtype=int),
|
|
)
|
|
annotator = PolygonAnnotator(color=Color.WHITE, color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(scene=scene.copy(), detections=detections)
|
|
np.testing.assert_array_equal(result, scene)
|
|
|
|
def test_annotate_with_all_false_compact_mask_unchanged(self):
|
|
"""All-False CompactMask crop (no contours) leaves scene pixels unchanged.
|
|
|
|
A detection whose mask is entirely False produces no polygons; the
|
|
annotator must not paint any pixels for that detection.
|
|
"""
|
|
height, width = 60, 80
|
|
scene = np.zeros((height, width, 3), dtype=np.uint8)
|
|
masks = np.zeros((1, height, width), dtype=bool)
|
|
xyxy = np.array([[10.0, 10.0, 50.0, 50.0]], dtype=np.float32)
|
|
detections = Detections(
|
|
xyxy=xyxy,
|
|
mask=CompactMask.from_dense(masks, xyxy, (height, width)),
|
|
class_id=np.array([0]),
|
|
)
|
|
annotator = PolygonAnnotator(color=Color.WHITE, color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(scene=scene.copy(), detections=detections)
|
|
np.testing.assert_array_equal(result, scene)
|
|
|
|
def test_annotate_with_single_compact_mask_detection(self):
|
|
"""N=1 CompactMask detection annotates the same as dense mask."""
|
|
height, width = 60, 80
|
|
scene = np.zeros((height, width, 3), dtype=np.uint8)
|
|
mask = np.zeros((height, width), dtype=bool)
|
|
mask[20:40, 15:55] = True
|
|
xyxy = np.array([[15.0, 20.0, 54.0, 39.0]], dtype=np.float32)
|
|
dense = Detections(xyxy=xyxy, mask=np.array([mask]), class_id=np.array([0]))
|
|
compact = Detections(
|
|
xyxy=xyxy,
|
|
mask=CompactMask.from_dense(np.array([mask]), xyxy, (height, width)),
|
|
class_id=np.array([0]),
|
|
)
|
|
annotator = PolygonAnnotator(
|
|
color=Color.WHITE, thickness=1, color_lookup=ColorLookup.INDEX
|
|
)
|
|
expected = annotator.annotate(scene=scene.copy(), detections=dense)
|
|
result = annotator.annotate(scene=scene.copy(), detections=compact)
|
|
np.testing.assert_array_equal(result, expected)
|
|
|
|
def test_annotate_compact_mask_float_xyxy_truncation(self):
|
|
"""Float xyxy with sub-pixel values are truncated to int before crop decode.
|
|
|
|
CompactMask stores bbox origins as int32 (via truncation); PolygonAnnotator
|
|
must produce the same output whether xyxy is integral or has fractional parts.
|
|
"""
|
|
height, width = 60, 80
|
|
scene = np.zeros((height, width, 3), dtype=np.uint8)
|
|
mask = np.zeros((height, width), dtype=bool)
|
|
mask[10:30, 5:45] = True
|
|
xyxy_int = np.array([[5.0, 10.0, 44.0, 29.0]], dtype=np.float32)
|
|
xyxy_float = np.array([[5.7, 10.9, 44.3, 29.1]], dtype=np.float32)
|
|
compact_int = Detections(
|
|
xyxy=xyxy_int,
|
|
mask=CompactMask.from_dense(np.array([mask]), xyxy_int, (height, width)),
|
|
class_id=np.array([0]),
|
|
)
|
|
compact_float = Detections(
|
|
xyxy=xyxy_float,
|
|
mask=CompactMask.from_dense(np.array([mask]), xyxy_float, (height, width)),
|
|
class_id=np.array([0]),
|
|
)
|
|
annotator = PolygonAnnotator(
|
|
color=Color.WHITE, thickness=1, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result_int = annotator.annotate(scene=scene.copy(), detections=compact_int)
|
|
result_float = annotator.annotate(scene=scene.copy(), detections=compact_float)
|
|
np.testing.assert_array_equal(result_int, result_float)
|
|
|
|
def test_compact_mask_disjoint_contours_offset_correct(self):
|
|
"""Disjoint-contour mask: both polygon blobs translate to correct image coords.
|
|
|
|
Verifies coordinate-level correctness of crop→image offset for a mask
|
|
with two separate blobs. After annotation, boundary pixels of each blob
|
|
must be painted; pixels between the blobs must remain unpainted.
|
|
"""
|
|
height, width = 80, 90
|
|
scene = np.zeros((height, width, 3), dtype=np.uint8)
|
|
mask = np.zeros((height, width), dtype=bool)
|
|
mask[10:20, 15:30] = True
|
|
mask[32:45, 40:55] = True
|
|
xyxy = np.array([[15.0, 10.0, 54.0, 44.0]], dtype=np.float32)
|
|
detections = Detections(
|
|
xyxy=xyxy,
|
|
mask=CompactMask.from_dense(np.array([mask]), xyxy, (height, width)),
|
|
class_id=np.array([0]),
|
|
)
|
|
annotator = PolygonAnnotator(
|
|
color=Color.WHITE, thickness=1, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=scene.copy(), detections=detections)
|
|
|
|
# Boundary of blob 1 in image space — must be painted
|
|
assert np.any(result[10:20, 15:30] != 0), "blob 1 boundary not painted"
|
|
# Boundary of blob 2 in image space — must be painted
|
|
assert np.any(result[32:45, 40:55] != 0), "blob 2 boundary not painted"
|
|
# Gap between blobs — no polygon pixels expected here
|
|
assert np.all(result[21:31, :] == 0), "gap between blobs has stray pixels"
|
|
|
|
|
|
class TestColorAnnotator:
|
|
"""Tests for ColorAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = ColorAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws a single color box"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = ColorAnnotator(
|
|
color=Color.RED, opacity=1.0, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.3)
|
|
|
|
|
|
class TestHaloAnnotator:
|
|
"""Tests for HaloAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = HaloAnnotator()
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_without_masks(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no masks"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = HaloAnnotator(color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_with_single_mask(self, test_image, test_mask):
|
|
"""Test that annotate method correctly draws a single halo"""
|
|
detections = _create_detections(
|
|
xyxy=[[10, 10, 90, 90]], mask=[test_mask], class_id=[0]
|
|
)
|
|
annotator = HaloAnnotator(
|
|
color=Color.BLUE,
|
|
opacity=0.8,
|
|
kernel_size=10,
|
|
color_lookup=ColorLookup.INDEX,
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.85)
|
|
|
|
def test_annotate_uint8_mask_matches_bool_mask(self, test_image, test_mask):
|
|
"""Test that uint8 and bool masks produce identical halos."""
|
|
detections_bool = _create_detections(
|
|
xyxy=[[10, 10, 90, 90]], mask=[test_mask], class_id=[0]
|
|
)
|
|
detections_uint8 = _create_detections(
|
|
xyxy=[[10, 10, 90, 90]], mask=[test_mask], class_id=[0]
|
|
)
|
|
detections_uint8.mask = detections_uint8.mask.astype(np.uint8)
|
|
|
|
annotator = HaloAnnotator(
|
|
color=Color.BLUE,
|
|
opacity=0.8,
|
|
kernel_size=10,
|
|
color_lookup=ColorLookup.INDEX,
|
|
)
|
|
result_bool = annotator.annotate(
|
|
scene=test_image.copy(), detections=detections_bool
|
|
)
|
|
result_uint8 = annotator.annotate(
|
|
scene=test_image.copy(), detections=detections_uint8
|
|
)
|
|
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])
|
|
|
|
def test_canvas_origin_nonzero_paints_at_correct_position(self):
|
|
"""Non-zero canvas_origin should offset mask coordinates into canvas."""
|
|
height, width = 30, 40
|
|
# Subcanvas covering region [10:25, 15:30] of the full image (15x15 px).
|
|
canvas = np.zeros((15, 15, 3), dtype=np.uint8)
|
|
# Mask in full-image coords: True at rows 12:18, cols 17:22.
|
|
# In canvas coords (origin=(15, 10)): rows 2:8, cols 2:7.
|
|
full_mask = np.zeros((height, width), dtype=bool)
|
|
full_mask[12:18, 17:22] = True
|
|
detections = Detections(
|
|
xyxy=np.array([[17.0, 12.0, 21.0, 17.0]]),
|
|
mask=full_mask[np.newaxis],
|
|
class_id=np.array([0]),
|
|
)
|
|
_paint_masks_by_area(
|
|
canvas,
|
|
detections,
|
|
Color.RED,
|
|
ColorLookup.INDEX,
|
|
canvas_origin=(15, 10),
|
|
)
|
|
# Pixels inside the mapped region should be BGR red (0, 0, 255)
|
|
assert np.all(canvas[2:8, 2:7] == (0, 0, 255))
|
|
# Pixels outside the painted region must remain zero
|
|
assert canvas[0, 0].tolist() == [0, 0, 0]
|
|
|
|
|
|
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: PolygonAnnotator(
|
|
color=Color.WHITE, thickness=1, color_lookup=ColorLookup.INDEX
|
|
),
|
|
id="polygon",
|
|
),
|
|
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"""
|
|
|
|
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())
|
|
|
|
def test_annotate_resets_when_resolution_changes(self) -> None:
|
|
"""Changing frame resolution must reset heat state instead of crashing."""
|
|
annotator = HeatMapAnnotator()
|
|
detections = _create_detections(xyxy=[[20, 20, 60, 60]])
|
|
first_scene = np.zeros((100, 100, 3), dtype=np.uint8)
|
|
second_scene = np.zeros((120, 80, 3), dtype=np.uint8)
|
|
|
|
annotator.annotate(scene=first_scene.copy(), detections=detections)
|
|
result = annotator.annotate(scene=second_scene.copy(), detections=detections)
|
|
|
|
assert result.shape == second_scene.shape
|
|
assert annotator.heat_mask is not None
|
|
assert annotator.heat_mask.shape == second_scene.shape[:2]
|
|
|
|
def test_annotate_hottest_region_survives_uint8_wrap(
|
|
self, test_image: np.ndarray
|
|
) -> None:
|
|
"""Heat count at 2^8=256 must not wrap uint8 to zero and blank the region."""
|
|
annotator = HeatMapAnnotator()
|
|
detections = _create_detections(xyxy=[[20, 20, 60, 60]])
|
|
for _ in range(256):
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
region_painted = np.count_nonzero(
|
|
np.any(result[20:60, 20:60] != test_image[20:60, 20:60], axis=2)
|
|
)
|
|
assert region_painted > 100
|
|
|
|
def test_reset_clears_accumulated_heat(self, test_image: np.ndarray) -> None:
|
|
"""reset() must zero accumulation so a reused annotator matches a fresh one.
|
|
|
|
The heatmap colours each pixel by its heat *relative to the current
|
|
maximum*, so a single uniformly-painted region always renders identically
|
|
regardless of its absolute count. To make the assertion actually depend on
|
|
reset having zeroed the buffer, heat is first built up on region A alone,
|
|
then after reset both region A and a fresh region B are annotated together.
|
|
If reset truly zeroed the buffer, A and B carry equal heat and the frame
|
|
matches a never-used annotator; if reset were a no-op, A's carried-over
|
|
count would dominate the max-normalisation and B would render a different
|
|
hue — so byte-equality with the fresh annotator can only hold when the
|
|
accumulation was genuinely discarded.
|
|
"""
|
|
region_a = _create_detections(xyxy=[[10, 10, 30, 30]])
|
|
region_a_and_b = _create_detections(xyxy=[[10, 10, 30, 30], [60, 60, 90, 90]])
|
|
reused = HeatMapAnnotator()
|
|
for _ in range(5):
|
|
reused.annotate(scene=test_image.copy(), detections=region_a)
|
|
reused.reset()
|
|
reused_result = reused.annotate(
|
|
scene=test_image.copy(), detections=region_a_and_b
|
|
)
|
|
|
|
fresh = HeatMapAnnotator()
|
|
fresh_result = fresh.annotate(
|
|
scene=test_image.copy(), detections=region_a_and_b
|
|
)
|
|
|
|
assert np.array_equal(reused_result, fresh_result)
|
|
|
|
|
|
class TestEllipseAnnotator:
|
|
"""Tests for EllipseAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = EllipseAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws a single ellipse"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = EllipseAnnotator(
|
|
color=Color.YELLOW, thickness=2, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.95)
|
|
|
|
|
|
class TestBoxCornerAnnotator:
|
|
"""Tests for BoxCornerAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = BoxCornerAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws box corners"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = BoxCornerAnnotator(
|
|
color=Color.WHITE,
|
|
thickness=3,
|
|
corner_length=10,
|
|
color_lookup=ColorLookup.INDEX,
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.95)
|
|
|
|
|
|
class TestCircleAnnotator:
|
|
"""Tests for CircleAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = CircleAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws a circle"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = CircleAnnotator(
|
|
color=Color.GREEN, thickness=2, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.95)
|
|
|
|
|
|
class TestDotAnnotator:
|
|
"""Tests for DotAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = DotAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws a dot"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = DotAnnotator(
|
|
color=Color.RED,
|
|
radius=5,
|
|
position=Position.CENTER,
|
|
color_lookup=ColorLookup.INDEX,
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.95)
|
|
|
|
|
|
class TestLabelAnnotator:
|
|
"""Tests for LabelAnnotator class"""
|
|
|
|
@pytest.mark.parametrize(
|
|
"border_radius",
|
|
[
|
|
pytest.param(0, id="radius-zero"),
|
|
pytest.param(-3, id="radius-negative"),
|
|
],
|
|
)
|
|
def test_draw_rounded_rectangle_square_matches_plain_rectangle(
|
|
self, border_radius: int
|
|
) -> None:
|
|
"""Non-positive radius fills the same pixels as a plain rectangle.
|
|
|
|
For border_radius < 0: previously raised cv2.error: radius >= 0 in
|
|
function 'circle'; fast path now silently draws square corners instead.
|
|
"""
|
|
scene = np.full((100, 120, 3), 9, dtype=np.uint8)
|
|
|
|
result = LabelAnnotator.draw_rounded_rectangle(
|
|
scene=scene.copy(),
|
|
xyxy=(10, 20, 90, 70),
|
|
color=(0, 0, 255),
|
|
border_radius=border_radius,
|
|
)
|
|
|
|
expected = scene.copy()
|
|
expected[20:71, 10:91] = (0, 0, 255)
|
|
assert np.array_equal(result, expected)
|
|
|
|
def test_draw_rounded_rectangle_clamped_to_zero_acts_as_square(self) -> None:
|
|
"""Positive border_radius clamped to 0 by a degenerate box draws square corners.
|
|
|
|
1px-wide box: min(10, 1 // 2) = min(10, 0) = 0 → fast path fires even
|
|
though the caller passed a positive radius.
|
|
"""
|
|
scene = np.full((100, 120, 3), 9, dtype=np.uint8)
|
|
|
|
result = LabelAnnotator.draw_rounded_rectangle(
|
|
scene=scene.copy(),
|
|
xyxy=(10, 20, 11, 70),
|
|
color=(0, 0, 255),
|
|
border_radius=10,
|
|
)
|
|
|
|
expected = scene.copy()
|
|
expected[20:71, 10:12] = (0, 0, 255)
|
|
assert np.array_equal(result, expected)
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = LabelAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws a label"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = LabelAnnotator(color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(
|
|
scene=test_image.copy(), detections=detections, labels=["test"]
|
|
)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.93)
|
|
|
|
def test_smart_position_spreads_boxes_once(
|
|
self, monkeypatch: pytest.MonkeyPatch, test_image: np.ndarray
|
|
) -> None:
|
|
"""smart_position should spread labels once per annotate call."""
|
|
calls = 0
|
|
original_spread_out_boxes = annotators_core.spread_out_boxes
|
|
|
|
def counting_spread_out_boxes(
|
|
boxes: np.ndarray, *args: object, **kwargs: object
|
|
) -> np.ndarray:
|
|
nonlocal calls
|
|
calls += 1
|
|
return original_spread_out_boxes(boxes, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(
|
|
annotators_core, "spread_out_boxes", counting_spread_out_boxes
|
|
)
|
|
|
|
detections = _create_detections(
|
|
xyxy=[[10, 10, 90, 90], [15, 15, 85, 85]], class_id=[0, 1]
|
|
)
|
|
annotator = LabelAnnotator(color_lookup=ColorLookup.INDEX, smart_position=True)
|
|
|
|
annotator.annotate(
|
|
scene=test_image.copy(), detections=detections, labels=["one", "two"]
|
|
)
|
|
|
|
assert calls == 1
|
|
|
|
|
|
class TestRichLabelAnnotator:
|
|
"""Tests for RichLabelAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = RichLabelAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws a rich label"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = RichLabelAnnotator(color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(
|
|
scene=test_image.copy(), detections=detections, labels=["test"]
|
|
)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.95)
|
|
|
|
def test_smart_position_spreads_boxes_once(
|
|
self, monkeypatch: pytest.MonkeyPatch, test_image: np.ndarray
|
|
) -> None:
|
|
"""smart_position should spread rich labels once per annotate call."""
|
|
calls = 0
|
|
original_spread_out_boxes = annotators_core.spread_out_boxes
|
|
|
|
def counting_spread_out_boxes(
|
|
boxes: np.ndarray, *args: object, **kwargs: object
|
|
) -> np.ndarray:
|
|
nonlocal calls
|
|
calls += 1
|
|
return original_spread_out_boxes(boxes, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(
|
|
annotators_core, "spread_out_boxes", counting_spread_out_boxes
|
|
)
|
|
|
|
detections = _create_detections(
|
|
xyxy=[[10, 10, 90, 90], [15, 15, 85, 85]], class_id=[0, 1]
|
|
)
|
|
annotator = RichLabelAnnotator(
|
|
color_lookup=ColorLookup.INDEX, smart_position=True
|
|
)
|
|
|
|
annotator.annotate(
|
|
scene=Image.fromarray(test_image.copy()),
|
|
detections=detections,
|
|
labels=["one", "two"],
|
|
)
|
|
|
|
assert calls == 1
|
|
|
|
|
|
class TestBlurAnnotator:
|
|
"""Tests for BlurAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = BlurAnnotator()
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_with_single_detection(self, gradient_image):
|
|
"""Test that annotate method correctly blurs a region"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = BlurAnnotator(kernel_size=15)
|
|
result = annotator.annotate(scene=gradient_image.copy(), detections=detections)
|
|
assert not np.array_equal(gradient_image, result)
|
|
|
|
@pytest.mark.parametrize("bad_size", [0, -1, -10])
|
|
def test_invalid_kernel_size_raises(self, bad_size):
|
|
"""BlurAnnotator must reject kernel_size < 1 at construction time."""
|
|
with pytest.raises(ValueError, match="kernel_size must be >= 1"):
|
|
BlurAnnotator(kernel_size=bad_size)
|
|
|
|
def test_annotate_zero_area_bbox_is_skipped(self, test_image):
|
|
"""Zero-area bounding boxes must be silently skipped, not crash."""
|
|
detections = _create_detections(xyxy=[[10, 10, 10, 50]], class_id=[0])
|
|
annotator = BlurAnnotator(kernel_size=5)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
|
|
class TestPixelateAnnotator:
|
|
"""Tests for PixelateAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = PixelateAnnotator()
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_with_single_detection(self, gradient_image):
|
|
"""Test that annotate method correctly pixelates a region"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = PixelateAnnotator(pixel_size=10)
|
|
result = annotator.annotate(scene=gradient_image.copy(), detections=detections)
|
|
assert not np.array_equal(gradient_image, result)
|
|
|
|
def test_annotate_bbox_smaller_than_pixel_size_does_not_raise(self):
|
|
"""PixelateAnnotator must not crash when the bbox is smaller than pixel_size.
|
|
|
|
Regression test for https://github.com/roboflow/supervision/issues/703:
|
|
a fixed pixel_size larger than the detection dimensions previously caused
|
|
an OpenCV assertion error in cv2.resize.
|
|
"""
|
|
image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
|
|
# bbox is 5x5; pixel_size=50 is much larger, triggers the avg-fill fallback
|
|
detections = _create_detections(xyxy=[[10, 10, 15, 15]], class_id=[0])
|
|
annotator = PixelateAnnotator(pixel_size=50)
|
|
result = annotator.annotate(scene=image.copy(), detections=detections)
|
|
assert result.shape == image.shape
|
|
|
|
def test_annotate_grayscale_image_does_not_raise(self):
|
|
"""PixelateAnnotator must work on single-channel (grayscale) images.
|
|
|
|
The small-ROI avg-fill branch previously sliced cv2.mean()[:3] into a
|
|
2-D array, causing a NumPy broadcast error on grayscale frames.
|
|
"""
|
|
gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
|
|
# Normal-size detection — exercises the resize path on a grayscale frame
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = PixelateAnnotator(pixel_size=10)
|
|
result = annotator.annotate(scene=gray.copy(), detections=detections)
|
|
assert result.shape == gray.shape
|
|
|
|
def test_annotate_grayscale_image_small_roi_does_not_raise(self):
|
|
"""Grayscale image with bbox smaller than pixel_size uses scalar avg fill.
|
|
|
|
Exercises the ndim-aware branch added to the small-ROI fallback.
|
|
"""
|
|
gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
|
|
detections = _create_detections(xyxy=[[10, 10, 15, 15]], class_id=[0])
|
|
annotator = PixelateAnnotator(pixel_size=50)
|
|
result = annotator.annotate(scene=gray.copy(), detections=detections)
|
|
assert result.shape == gray.shape
|
|
|
|
@pytest.mark.parametrize("bad_size", [0, -1, -10])
|
|
def test_invalid_pixel_size_raises(self, bad_size):
|
|
"""PixelateAnnotator must reject pixel_size < 1 at construction time."""
|
|
with pytest.raises(ValueError, match="pixel_size must be >= 1"):
|
|
PixelateAnnotator(pixel_size=bad_size)
|
|
|
|
def test_annotate_zero_area_bbox_is_skipped(self, test_image):
|
|
"""Zero-area bounding boxes must be silently skipped, not crash."""
|
|
detections = _create_detections(xyxy=[[10, 10, 10, 50]], class_id=[0])
|
|
annotator = PixelateAnnotator(pixel_size=5)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
|
|
class TestTriangleAnnotator:
|
|
"""Tests for TriangleAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = TriangleAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws a triangle"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = TriangleAnnotator(
|
|
color=Color.RED, base=20, height=20, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.95)
|
|
|
|
|
|
class TestRoundBoxAnnotator:
|
|
"""Tests for RoundBoxAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = RoundBoxAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws a round box"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = RoundBoxAnnotator(
|
|
color=Color.BLUE, thickness=2, roundness=0.5, color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.9)
|
|
|
|
|
|
class TestPercentageBarAnnotator:
|
|
"""Tests for PercentageBarAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = PercentageBarAnnotator()
|
|
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):
|
|
"""Test that annotate method correctly draws a percentage bar"""
|
|
detections = _create_detections(
|
|
xyxy=[[10, 10, 90, 90]], confidence=[0.75], class_id=[0]
|
|
)
|
|
annotator = PercentageBarAnnotator(color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert_image_mostly_same(test_image, result, similarity_threshold=0.93)
|
|
|
|
|
|
class TestPositionHelpers:
|
|
"""Tests for helper methods that map `Position` to coordinates."""
|
|
|
|
@pytest.mark.parametrize(
|
|
("helper", "args"),
|
|
[
|
|
pytest.param(
|
|
PercentageBarAnnotator.calculate_border_coordinates,
|
|
((10, 10), (4, 4), cast(Position, "invalid")),
|
|
id="percentage-bar",
|
|
),
|
|
pytest.param(
|
|
CropAnnotator.calculate_crop_coordinates,
|
|
((10, 10), (4, 4), cast(Position, "invalid")),
|
|
id="crop",
|
|
),
|
|
],
|
|
)
|
|
def test_unknown_position_raises(
|
|
self, helper: Any, args: tuple[Any, Any, Any]
|
|
) -> None:
|
|
"""Unsupported positions must raise instead of returning None."""
|
|
with pytest.raises(ValueError, match="Unsupported position"):
|
|
helper(*args)
|
|
|
|
|
|
class TestCropAnnotator:
|
|
"""Tests for CropAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = CropAnnotator()
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_with_single_detection(self, gradient_image):
|
|
"""Test that annotate method correctly draws a crop"""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = CropAnnotator(border_color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(scene=gradient_image.copy(), detections=detections)
|
|
assert not np.array_equal(gradient_image, result)
|
|
|
|
def test_annotate_emits_no_deprecation_warning(self, gradient_image):
|
|
"""Internal overlay must not surface the deprecated `overlay_image` warning."""
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]], class_id=[0])
|
|
annotator = CropAnnotator(border_color_lookup=ColorLookup.INDEX)
|
|
with warnings.catch_warnings(record=True) as caught:
|
|
warnings.simplefilter("always")
|
|
annotator.annotate(scene=gradient_image.copy(), detections=detections)
|
|
deprecations = [
|
|
w
|
|
for w in caught
|
|
if issubclass(w.category, (DeprecationWarning, FutureWarning))
|
|
]
|
|
assert deprecations == []
|
|
|
|
def test_annotate_with_partially_out_of_bounds_detection(self, gradient_image):
|
|
"""Partially-OOB box is clipped and rendered; scene must change."""
|
|
detections = _create_detections(xyxy=[[-10, -10, 30, 30]], class_id=[0])
|
|
annotator = CropAnnotator(
|
|
position=Position.CENTER, border_color_lookup=ColorLookup.INDEX
|
|
)
|
|
result = annotator.annotate(scene=gradient_image.copy(), detections=detections)
|
|
assert result.shape == gradient_image.shape
|
|
assert not np.array_equal(gradient_image, result)
|
|
|
|
def test_annotate_with_fully_out_of_bounds_detection(self, gradient_image):
|
|
"""A box fully outside the scene collapses to zero area and is skipped."""
|
|
detections = _create_detections(xyxy=[[-50, -50, -10, -10]], class_id=[0])
|
|
annotator = CropAnnotator(border_color_lookup=ColorLookup.INDEX)
|
|
result = annotator.annotate(scene=gradient_image.copy(), detections=detections)
|
|
assert np.array_equal(gradient_image, result)
|
|
|
|
@pytest.mark.parametrize(
|
|
"xyxy",
|
|
[
|
|
pytest.param([-5, 20, 40, 60], id="negative-x-min"),
|
|
pytest.param([20, -5, 60, 40], id="negative-y-min"),
|
|
pytest.param([-10, -10, 30, 30], id="negative-x-and-y-min"),
|
|
pytest.param([60, 20, 140, 60], id="past-right-edge"),
|
|
pytest.param([-20, -20, 140, 140], id="larger-than-scene"),
|
|
],
|
|
)
|
|
def test_annotate_with_box_crossing_scene_border(
|
|
self, gradient_image, xyxy: list[int]
|
|
) -> None:
|
|
"""Boxes extending past the scene border are clipped instead of raising"""
|
|
detections = _create_detections(xyxy=[xyxy], class_id=[0])
|
|
annotator = CropAnnotator()
|
|
|
|
result = annotator.annotate(scene=gradient_image.copy(), detections=detections)
|
|
|
|
assert result.shape == gradient_image.shape
|
|
|
|
@pytest.mark.parametrize(
|
|
"xyxy",
|
|
[
|
|
pytest.param([150, 150, 200, 200], id="fully-outside"),
|
|
pytest.param([-50, -50, -10, -10], id="fully-negative"),
|
|
pytest.param([30, 20, 30, 60], id="zero-width"),
|
|
pytest.param([30, 30, 30, 30], id="zero-area"),
|
|
],
|
|
)
|
|
def test_annotate_skips_boxes_empty_after_clipping(
|
|
self, gradient_image, xyxy: list[int]
|
|
) -> None:
|
|
"""Boxes with no visible area are skipped instead of raising cv2.error"""
|
|
detections = _create_detections(xyxy=[xyxy], class_id=[0])
|
|
annotator = CropAnnotator()
|
|
|
|
result = annotator.annotate(scene=gradient_image.copy(), detections=detections)
|
|
|
|
assert np.array_equal(gradient_image, result)
|
|
|
|
def test_annotate_mixed_valid_and_degenerate_boxes(self, gradient_image) -> None:
|
|
"""A degenerate box does not prevent valid boxes from being drawn"""
|
|
detections = _create_detections(
|
|
xyxy=[[150, 150, 200, 200], [10, 10, 90, 90]], class_id=[0, 1]
|
|
)
|
|
annotator = CropAnnotator()
|
|
|
|
result = annotator.annotate(scene=gradient_image.copy(), detections=detections)
|
|
|
|
assert not np.array_equal(gradient_image, result)
|
|
|
|
def test_annotate_overlapping_crops_sample_from_original_scene(self) -> None:
|
|
"""Later crops must sample the original un-annotated scene.
|
|
|
|
box1 is pasted into the region that box2 crops from. Without the
|
|
source_scene = scene.copy() fix, box2 reads box1's paste value
|
|
instead of the original pixel — the aliasing regression.
|
|
"""
|
|
# Arrange: two distinct pixel bands; box1's paste region overlaps box2's crop
|
|
scene = np.full((80, 80, 3), 50, dtype=np.uint8)
|
|
scene[0:20, 0:20] = 10 # band A — box1 crops here (value 10)
|
|
scene[20:40, 20:40] = 200 # band B — box2 crops here; box1 pastes here
|
|
|
|
detections = _create_detections(
|
|
xyxy=[[0, 0, 20, 20], [20, 20, 40, 40]], class_id=[0, 1]
|
|
)
|
|
# BOTTOM_RIGHT: each crop is pasted at its (x2, y2) corner.
|
|
# box1 pastes band A (value 10) at rows 20-39, cols 20-39 — exactly
|
|
# where box2 will crop.
|
|
annotator = CropAnnotator(
|
|
position=Position.BOTTOM_RIGHT,
|
|
scale_factor=1.0,
|
|
)
|
|
|
|
# Act
|
|
result = annotator.annotate(scene=scene.copy(), detections=detections)
|
|
|
|
# Assert: box2's crop is pasted at rows 40-59, cols 40-59.
|
|
# Interior (rows 42-57, cols 42-57) avoids the 2-pixel default border
|
|
# and must equal 200 — the original band B value. Without the fix,
|
|
# box2 samples 10 from the painted scene instead of 200 from original.
|
|
interior = result[42:58, 42:58]
|
|
assert np.all(interior == 200), (
|
|
f"box2 crop paste region should contain original pixel 200, "
|
|
f"got {np.unique(interior).tolist()!r}. "
|
|
"Regression: source_scene must be scene.copy(), not an alias."
|
|
)
|
|
|
|
|
|
class TestIconAnnotator:
|
|
"""Tests for IconAnnotator class"""
|
|
|
|
def test_annotate_emits_no_deprecation_warning(self, test_image, tmp_path):
|
|
"""Internal overlay must not surface the deprecated `overlay_image` warning."""
|
|
icon_path = str(tmp_path / "icon.png")
|
|
icon = np.full((20, 20, 4), (0, 255, 0, 255), dtype=np.uint8)
|
|
cv2.imwrite(icon_path, icon)
|
|
detections = _create_detections(xyxy=[[20, 20, 60, 60]], class_id=[0])
|
|
annotator = IconAnnotator()
|
|
with warnings.catch_warnings(record=True) as caught:
|
|
warnings.simplefilter("always")
|
|
annotator.annotate(
|
|
scene=test_image.copy(), detections=detections, icon_path=icon_path
|
|
)
|
|
deprecations = [
|
|
w
|
|
for w in caught
|
|
if issubclass(w.category, (DeprecationWarning, FutureWarning))
|
|
]
|
|
assert deprecations == []
|
|
|
|
def test_icon_cache_is_shared_by_path_and_resolution(
|
|
self, monkeypatch, test_image, tmp_path
|
|
):
|
|
"""Equal path/resolution icon loads are cached across annotator instances."""
|
|
icon_path = str(tmp_path / "icon.png")
|
|
icon = np.full((20, 20, 4), (0, 255, 0, 255), dtype=np.uint8)
|
|
cv2.imwrite(icon_path, icon)
|
|
detections = _create_detections(xyxy=[[20, 20, 60, 60]], class_id=[0])
|
|
imread_calls = 0
|
|
original_imread = cv2.imread
|
|
|
|
def count_imread(path, flags):
|
|
nonlocal imread_calls
|
|
imread_calls += 1
|
|
return original_imread(path, flags)
|
|
|
|
monkeypatch.setattr(cv2, "imread", count_imread)
|
|
|
|
for _ in range(2):
|
|
IconAnnotator(icon_resolution_wh=(16, 16)).annotate(
|
|
scene=test_image.copy(), detections=detections, icon_path=icon_path
|
|
)
|
|
|
|
assert imread_calls == 1
|
|
|
|
|
|
class TestBackgroundOverlayAnnotator:
|
|
"""Tests for BackgroundOverlayAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections = Detections.empty()
|
|
annotator = BackgroundOverlayAnnotator()
|
|
result = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_with_single_detection(self):
|
|
"""Test that annotate method correctly draws background overlay"""
|
|
image = np.ones((100, 100, 3), dtype=np.uint8) * 255
|
|
detections = _create_detections(xyxy=[[10, 10, 90, 90]])
|
|
annotator = BackgroundOverlayAnnotator(color=Color.BLACK, opacity=0.5)
|
|
result = annotator.annotate(scene=image.copy(), detections=detections)
|
|
assert not np.array_equal(image, result)
|
|
|
|
@pytest.mark.parametrize(
|
|
("xyxy", "inside_xy", "outside_xy"),
|
|
[
|
|
pytest.param([-5, 20, 40, 60], (20, 30), (60, 80), id="crosses-left-edge"),
|
|
pytest.param([20, -5, 60, 40], (30, 20), (80, 60), id="crosses-top-edge"),
|
|
pytest.param(
|
|
[-10, -10, 40, 40], (20, 20), (70, 70), id="crosses-both-edges"
|
|
),
|
|
],
|
|
)
|
|
def test_annotate_preserves_detection_crossing_scene_border(
|
|
self, xyxy: list[int], inside_xy: tuple[int, int], outside_xy: tuple[int, int]
|
|
) -> None:
|
|
"""The visible part of a box crossing the border keeps original pixels"""
|
|
image = np.full((100, 100, 3), 200, dtype=np.uint8)
|
|
detections = _create_detections(xyxy=[xyxy])
|
|
annotator = BackgroundOverlayAnnotator(color=Color.BLACK, opacity=0.5)
|
|
|
|
result = annotator.annotate(scene=image.copy(), detections=detections)
|
|
|
|
x_in, y_in = inside_xy
|
|
x_out, y_out = outside_xy
|
|
assert np.array_equal(result[y_in, x_in], np.array([200, 200, 200]))
|
|
assert np.array_equal(result[y_out, x_out], np.array([100, 100, 100]))
|
|
|
|
def test_annotate_fully_out_negative_box_does_not_corrupt(self) -> None:
|
|
"""Both-negative OOB box must not restore an in-bounds region via wrap-around"""
|
|
image = np.full((100, 100, 3), 200, dtype=np.uint8)
|
|
detections = _create_detections(xyxy=[[-30, -30, -5, -5]])
|
|
annotator = BackgroundOverlayAnnotator(color=Color.BLACK, opacity=0.5)
|
|
|
|
result = annotator.annotate(scene=image.copy(), detections=detections)
|
|
|
|
assert np.array_equal(result[80, 80], np.array([100, 100, 100]))
|
|
|
|
def test_annotate_force_box_preserves_detection_crossing_scene_border(self):
|
|
"""force_box with a border-crossing box keeps the visible detection region"""
|
|
image = np.full((100, 100, 3), 200, dtype=np.uint8)
|
|
mask = np.zeros((100, 100), dtype=bool)
|
|
mask[20:60, 0:40] = True
|
|
detections = _create_detections(xyxy=[[-5, 20, 40, 60]], mask=[mask])
|
|
annotator = BackgroundOverlayAnnotator(
|
|
color=Color.BLACK, opacity=0.5, force_box=True
|
|
)
|
|
|
|
result = annotator.annotate(scene=image.copy(), detections=detections)
|
|
|
|
assert np.array_equal(result[30, 20], np.array([200, 200, 200]))
|
|
assert np.array_equal(result[80, 60], np.array([100, 100, 100]))
|
|
|
|
def test_annotate_with_fully_out_of_bounds_detection(self):
|
|
"""A box fully outside the scene leaves the whole scene tinted"""
|
|
image = np.full((100, 100, 3), 200, dtype=np.uint8)
|
|
detections = _create_detections(xyxy=[[150, 150, 200, 200]])
|
|
annotator = BackgroundOverlayAnnotator(color=Color.BLACK, opacity=0.5)
|
|
|
|
result = annotator.annotate(scene=image.copy(), detections=detections)
|
|
|
|
assert np.all(result == 100)
|
|
|
|
def test_annotate_uint8_mask_matches_bool_mask(self):
|
|
"""Test that uint8 and bool masks produce identical overlays."""
|
|
image = np.ones((100, 100, 3), dtype=np.uint8) * 255
|
|
mask = np.zeros((100, 100), dtype=bool)
|
|
mask[10:90, 10:90] = True
|
|
|
|
detections_bool = _create_detections(xyxy=[[10, 10, 90, 90]], mask=[mask])
|
|
detections_uint8 = _create_detections(xyxy=[[10, 10, 90, 90]], mask=[mask])
|
|
detections_uint8.mask = detections_uint8.mask.astype(np.uint8)
|
|
|
|
annotator = BackgroundOverlayAnnotator(color=Color.BLACK, opacity=0.5)
|
|
result_bool = annotator.annotate(scene=image.copy(), detections=detections_bool)
|
|
result_uint8 = annotator.annotate(
|
|
scene=image.copy(), detections=detections_uint8
|
|
)
|
|
assert np.array_equal(result_bool, result_uint8)
|
|
|
|
|
|
class TestComparisonAnnotator:
|
|
"""Tests for ComparisonAnnotator class"""
|
|
|
|
def test_annotate_with_no_detections(self, test_image):
|
|
"""Test that annotate method returns unmodified image when no detections"""
|
|
detections1 = Detections.empty()
|
|
detections2 = Detections.empty()
|
|
annotator = ComparisonAnnotator()
|
|
result = annotator.annotate(
|
|
scene=test_image.copy(), detections_1=detections1, detections_2=detections2
|
|
)
|
|
assert np.array_equal(test_image, result)
|
|
|
|
def test_annotate_with_single_detection_each(self):
|
|
"""Test that annotate method correctly compares two detections"""
|
|
image = np.ones((100, 100, 3), dtype=np.uint8) * 255
|
|
detections1 = _create_detections(xyxy=[[10, 10, 50, 50]])
|
|
detections2 = _create_detections(xyxy=[[30, 30, 70, 70]])
|
|
annotator = ComparisonAnnotator()
|
|
result = annotator.annotate(
|
|
scene=image.copy(), detections_1=detections1, detections_2=detections2
|
|
)
|
|
assert not np.array_equal(image, result)
|
|
|
|
|
|
class TestTraceAnnotatorReset:
|
|
"""Tests for TraceAnnotator.reset() clearing accumulated trace history."""
|
|
|
|
def test_reset_empties_trace_buffers(self, test_image: np.ndarray) -> None:
|
|
"""reset() must clear the underlying Trace buffers to their empty state."""
|
|
annotator = TraceAnnotator(trace_length=10)
|
|
detections = _create_detections(
|
|
xyxy=[[10, 10, 30, 30]], class_id=[1], tracker_id=[7]
|
|
)
|
|
annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
|
|
annotator.reset()
|
|
|
|
assert annotator.trace.frame_id.shape == (0,)
|
|
assert annotator.trace.xy.shape == (0, 2)
|
|
assert annotator.trace.tracker_id.shape == (0,)
|
|
assert annotator.trace.current_frame_id == 0
|
|
|
|
def test_reset_matches_fresh_annotator(self, test_image: np.ndarray) -> None:
|
|
"""After reset() a reused annotator must render identically to a fresh one.
|
|
|
|
The two streams reuse the same ``tracker_id`` but follow spatially
|
|
distinct paths. If reset were a no-op, ``Trace.get`` would return the
|
|
first stream's points concatenated with the second's and draw a spurious
|
|
polyline bridging the two paths; only a genuine reset leaves solely the
|
|
second stream's points, so byte-equality with a never-used annotator can
|
|
hold only when the prior history was actually discarded. ``trace_length``
|
|
is large enough that no windowing prunes away the stale points that a
|
|
broken reset would leave behind.
|
|
"""
|
|
first_stream = [
|
|
_create_detections(
|
|
xyxy=[[10 + step * 6, 10 + step * 6, 20 + step * 6, 20 + step * 6]],
|
|
class_id=[1],
|
|
tracker_id=[7],
|
|
)
|
|
for step in range(5)
|
|
]
|
|
second_stream = [
|
|
_create_detections(
|
|
xyxy=[[80 - step * 6, 10 + step * 6, 90 - step * 6, 20 + step * 6]],
|
|
class_id=[1],
|
|
tracker_id=[7],
|
|
)
|
|
for step in range(5)
|
|
]
|
|
reused = TraceAnnotator(trace_length=30)
|
|
reused_scene = test_image.copy()
|
|
for detections in first_stream:
|
|
reused_scene = reused.annotate(scene=reused_scene, detections=detections)
|
|
reused.reset()
|
|
reused_scene = test_image.copy()
|
|
for detections in second_stream:
|
|
reused_scene = reused.annotate(scene=reused_scene, detections=detections)
|
|
|
|
fresh = TraceAnnotator(trace_length=30)
|
|
fresh_scene = test_image.copy()
|
|
for detections in second_stream:
|
|
fresh_scene = fresh.annotate(scene=fresh_scene, detections=detections)
|
|
|
|
assert np.array_equal(reused_scene, fresh_scene)
|
|
|
|
|
|
class TestTraceAnnotatorSmoothStationary:
|
|
"""Regression tests for TraceAnnotator(smooth=True) on stationary tracker ids."""
|
|
|
|
def test_stationary_tracker_does_not_crash_spline_fit(self, test_image):
|
|
"""
|
|
When the same tracker stays at an identical anchor point for several
|
|
frames the trace buffer accumulates duplicate points. `scipy.splprep`
|
|
rejects a zero-length input curve with `ValueError: Invalid inputs.`,
|
|
so the annotator must survive this input without raising.
|
|
"""
|
|
detections = _create_detections(
|
|
xyxy=[[100, 100, 120, 120]],
|
|
class_id=[1],
|
|
tracker_id=[42],
|
|
)
|
|
annotator = TraceAnnotator(smooth=True, trace_length=10)
|
|
scene = test_image.copy()
|
|
for _ in range(6):
|
|
scene = annotator.annotate(scene=scene, detections=detections)
|
|
assert scene.shape == test_image.shape
|
|
|
|
def test_smooth_trace_still_renders_for_moving_tracker(self, test_image):
|
|
"""Moving tracker must produce a spline trace distinct from the raw polyline.
|
|
|
|
Compares smooth=True output against smooth=False for the same movement
|
|
path to confirm the smoothing path is actually exercised (not just that
|
|
some pixels changed).
|
|
"""
|
|
smooth_annotator = TraceAnnotator(smooth=True, trace_length=10, thickness=2)
|
|
raw_annotator = TraceAnnotator(smooth=False, trace_length=10, thickness=2)
|
|
scene_smooth = test_image.copy()
|
|
scene_raw = test_image.copy()
|
|
for offset in range(6):
|
|
detections = _create_detections(
|
|
xyxy=[
|
|
[10 + offset * 5, 10 + offset * 5, 30 + offset * 5, 30 + offset * 5]
|
|
],
|
|
class_id=[1],
|
|
tracker_id=[7],
|
|
)
|
|
scene_smooth = smooth_annotator.annotate(
|
|
scene=scene_smooth, detections=detections
|
|
)
|
|
scene_raw = raw_annotator.annotate(scene=scene_raw, detections=detections)
|
|
# After 4+ unique anchor positions the spline path fires and diverges from the
|
|
# raw polyline — the two output images must differ.
|
|
assert not np.array_equal(scene_smooth, scene_raw)
|
|
|
|
@pytest.mark.parametrize(
|
|
"unique_positions",
|
|
[1, 2, 3, 4],
|
|
ids=["1_unique", "2_unique", "3_unique", "4_unique"],
|
|
)
|
|
def test_smooth_does_not_crash_for_unique_point_counts(
|
|
self, test_image, unique_positions
|
|
):
|
|
"""smooth=True must not crash for any unique-position count from 1 to 4.
|
|
|
|
Each position is repeated twice to simulate brief holds between moves.
|
|
Covers the boundary at len(unique_xy) == 4 where splprep first fires.
|
|
"""
|
|
annotator = TraceAnnotator(smooth=True, trace_length=10, thickness=2)
|
|
scene = test_image.copy()
|
|
for pos_idx in range(unique_positions):
|
|
for _ in range(2):
|
|
x = 10 + pos_idx * 15
|
|
detections = _create_detections(
|
|
xyxy=[[x, x, x + 15, x + 15]],
|
|
class_id=[1],
|
|
tracker_id=[99],
|
|
)
|
|
scene = annotator.annotate(scene=scene, detections=detections)
|
|
assert scene.shape == test_image.shape
|
|
|
|
def test_smooth_fallback_matches_raw_when_fewer_than_four_unique_points(
|
|
self, test_image
|
|
):
|
|
"""With <4 unique positions smooth=True output must match smooth=False.
|
|
|
|
Verifies the dedup-then-fallback path: when unique_xy has ≤3 points,
|
|
both branches use the same raw-polyline draw.
|
|
"""
|
|
annotator_smooth = TraceAnnotator(smooth=True, trace_length=10, thickness=2)
|
|
annotator_raw = TraceAnnotator(smooth=False, trace_length=10, thickness=2)
|
|
scene_smooth = test_image.copy()
|
|
scene_raw = test_image.copy()
|
|
for pos_idx in range(3):
|
|
for _ in range(2):
|
|
x = 10 + pos_idx * 15
|
|
detections = _create_detections(
|
|
xyxy=[[x, x, x + 15, x + 15]],
|
|
class_id=[1],
|
|
tracker_id=[99],
|
|
)
|
|
scene_smooth = annotator_smooth.annotate(
|
|
scene=scene_smooth, detections=detections
|
|
)
|
|
scene_raw = annotator_raw.annotate(
|
|
scene=scene_raw, detections=detections
|
|
)
|
|
assert np.array_equal(scene_smooth, scene_raw)
|
|
|
|
def test_smooth_true_single_frame_does_not_crash(self, test_image):
|
|
"""A single annotate() call with smooth=True must not crash.
|
|
|
|
When len(xy) == 1 the drawing guard skips cv2.polylines entirely;
|
|
the dedup path runs safely on an empty np.diff result.
|
|
"""
|
|
detections = _create_detections(
|
|
xyxy=[[50, 50, 70, 70]],
|
|
class_id=[1],
|
|
tracker_id=[1],
|
|
)
|
|
annotator = TraceAnnotator(smooth=True, trace_length=10)
|
|
scene = annotator.annotate(scene=test_image.copy(), detections=detections)
|
|
assert scene.shape == test_image.shape
|
|
|
|
def test_smooth_false_stationary_tracker_does_not_crash(self, test_image):
|
|
"""smooth=False with a stationary tracker must not crash (regression guard).
|
|
|
|
Ensures the refactor did not accidentally alter the smooth=False code path.
|
|
"""
|
|
detections = _create_detections(
|
|
xyxy=[[100, 100, 120, 120]],
|
|
class_id=[1],
|
|
tracker_id=[42],
|
|
)
|
|
annotator = TraceAnnotator(smooth=False, trace_length=10)
|
|
scene = test_image.copy()
|
|
for _ in range(6):
|
|
scene = annotator.annotate(scene=scene, detections=detections)
|
|
assert scene.shape == test_image.shape
|