374 lines
14 KiB
Python
374 lines
14 KiB
Python
"""Tests for src/supervision/detection/tools/transformers.py processing functions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from supervision.config import CLASS_NAME_DATA_FIELD
|
|
from supervision.detection.tools.transformers import (
|
|
append_class_names_to_data,
|
|
png_string_to_segmentation_array,
|
|
process_transformers_detection_result,
|
|
process_transformers_v4_panoptic_segmentation_result,
|
|
process_transformers_v4_segmentation_result,
|
|
process_transformers_v5_panoptic_segmentation_result,
|
|
process_transformers_v5_segmentation_result,
|
|
process_transformers_v5_semantic_or_instance_segmentation_result,
|
|
)
|
|
from tests.helpers import _FakeDetachTensor, make_panoptic_png
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# png_string_to_segmentation_array
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPngStringToSegmentationArray:
|
|
"""png_string_to_segmentation_array extracts the red channel as a label map."""
|
|
|
|
def test_extracts_red_channel_as_segment_ids(self) -> None:
|
|
"""RGBA PNG: red channel values become the returned label array."""
|
|
seg_map = np.array([[1, 2], [3, 0]], dtype=np.uint8)
|
|
png_bytes = make_panoptic_png(seg_map)
|
|
|
|
result = png_string_to_segmentation_array(png_bytes)
|
|
|
|
np.testing.assert_array_equal(result, seg_map)
|
|
|
|
def test_returns_array_of_shape_h_w(self) -> None:
|
|
"""Output shape matches the image height and width."""
|
|
seg_map = np.zeros((6, 8), dtype=np.uint8)
|
|
seg_map[2:4, 3:5] = 7
|
|
png_bytes = make_panoptic_png(seg_map)
|
|
|
|
result = png_string_to_segmentation_array(png_bytes)
|
|
|
|
assert result.shape == (6, 8)
|
|
assert result[2, 3] == 7
|
|
assert result[0, 0] == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# append_class_names_to_data
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAppendClassNamesToData:
|
|
"""append_class_names_to_data conditionally populates CLASS_NAME_DATA_FIELD."""
|
|
|
|
def test_with_id2label_adds_class_names_array(self) -> None:
|
|
"""When id2label provided, CLASS_NAME_DATA_FIELD is set to mapped names."""
|
|
class_ids = np.array([0, 1, 0])
|
|
id2label = {0: "cat", 1: "dog"}
|
|
|
|
result = append_class_names_to_data(class_ids, id2label, {})
|
|
|
|
np.testing.assert_array_equal(
|
|
result[CLASS_NAME_DATA_FIELD], ["cat", "dog", "cat"]
|
|
)
|
|
|
|
def test_without_id2label_returns_unchanged_data(self) -> None:
|
|
"""When id2label is None, no class name key is written."""
|
|
class_ids = np.array([0, 1])
|
|
|
|
result = append_class_names_to_data(class_ids, None, {})
|
|
|
|
assert CLASS_NAME_DATA_FIELD not in result
|
|
|
|
def test_merges_into_existing_data_dict(self) -> None:
|
|
"""Existing data dict keys are preserved when class names are added."""
|
|
existing = {"custom_key": np.array([1, 2])}
|
|
|
|
result = append_class_names_to_data(np.array([0]), {0: "cat"}, existing)
|
|
|
|
assert "custom_key" in result
|
|
assert CLASS_NAME_DATA_FIELD in result
|
|
|
|
def test_empty_class_ids_with_id2label_yields_empty_name_array(self) -> None:
|
|
"""Zero detections with id2label still produce an empty name array."""
|
|
result = append_class_names_to_data(np.array([]), {0: "cat"}, {})
|
|
|
|
assert CLASS_NAME_DATA_FIELD in result
|
|
assert len(result[CLASS_NAME_DATA_FIELD]) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# process_transformers_detection_result
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestProcessTransformersDetectionResult:
|
|
"""process_transformers_detection_result extracts xyxy/confidence/class_id."""
|
|
|
|
@pytest.mark.parametrize(
|
|
("n_boxes", "with_id2label"),
|
|
[
|
|
pytest.param(2, False, id="two-detections-no-labels"),
|
|
pytest.param(1, True, id="single-detection-with-labels"),
|
|
pytest.param(0, False, id="empty-no-labels"),
|
|
],
|
|
)
|
|
def test_maps_fields_and_optionally_class_names(
|
|
self, n_boxes: int, with_id2label: bool
|
|
) -> None:
|
|
"""Output has xyxy, confidence, class_id; class names when id2label given."""
|
|
xyxy = np.zeros((n_boxes, 4), dtype=np.float32)
|
|
scores = np.ones(n_boxes, dtype=np.float32) * 0.9
|
|
labels = np.arange(n_boxes, dtype=np.int64)
|
|
id2label = {i: f"cls{i}" for i in range(n_boxes)} if with_id2label else None
|
|
detection_result = {
|
|
"boxes": _FakeDetachTensor(xyxy),
|
|
"scores": _FakeDetachTensor(scores),
|
|
"labels": _FakeDetachTensor(labels),
|
|
}
|
|
|
|
out = process_transformers_detection_result(detection_result, id2label)
|
|
|
|
assert out["xyxy"].shape == (n_boxes, 4)
|
|
assert len(out["confidence"]) == n_boxes
|
|
assert len(out["class_id"]) == n_boxes
|
|
if with_id2label and n_boxes > 0:
|
|
assert CLASS_NAME_DATA_FIELD in out["data"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# process_transformers_v4_segmentation_result
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestProcessTransformersV4SegmentationResult:
|
|
"""process_transformers_v4_segmentation_result handles masks, boxes, panoptic."""
|
|
|
|
def test_masks_only_path_uses_mask_to_xyxy(self) -> None:
|
|
"""Without boxes, mask_to_xyxy derives xyxy; mask shape is (N, H, W)."""
|
|
masks = np.zeros((2, 4, 4), dtype=bool)
|
|
masks[0, 0:2, 0:2] = True
|
|
masks[1, 2:4, 2:4] = True
|
|
seg_result = {
|
|
"masks": _FakeDetachTensor(masks.astype(np.uint8)),
|
|
"labels": _FakeDetachTensor(np.array([0, 1], dtype=np.int64)),
|
|
"scores": _FakeDetachTensor(np.array([0.9, 0.8], dtype=np.float32)),
|
|
}
|
|
|
|
out = process_transformers_v4_segmentation_result(seg_result, None)
|
|
|
|
assert out["mask"].shape == (2, 4, 4)
|
|
assert out["xyxy"].shape == (2, 4)
|
|
|
|
def test_masks_with_boxes_squeezes_mask_axis(self) -> None:
|
|
"""When boxes provided, masks (N,1,H,W) are squeezed to (N,H,W)."""
|
|
masks = np.zeros((1, 1, 4, 4), dtype=bool)
|
|
masks[0, 0, 0:2, 0:2] = True
|
|
seg_result = {
|
|
"boxes": _FakeDetachTensor(np.array([[0, 0, 2, 2]], dtype=np.float32)),
|
|
"masks": _FakeDetachTensor(masks.astype(np.uint8)),
|
|
"labels": _FakeDetachTensor(np.array([3], dtype=np.int64)),
|
|
"scores": _FakeDetachTensor(np.array([0.75], dtype=np.float32)),
|
|
}
|
|
|
|
out = process_transformers_v4_segmentation_result(seg_result, None)
|
|
|
|
assert out["mask"].shape == (1, 4, 4)
|
|
|
|
def test_panoptic_path_triggered_by_png_string(self) -> None:
|
|
"""png_string key routes to panoptic sub-processor; returns mask per segment."""
|
|
seg_map = np.zeros((4, 4), dtype=np.uint8)
|
|
seg_map[0:2, 0:2] = 1
|
|
seg_result = {
|
|
"png_string": make_panoptic_png(seg_map),
|
|
"segments_info": [{"id": 1, "category_id": 5}],
|
|
}
|
|
|
|
out = process_transformers_v4_segmentation_result(seg_result, None)
|
|
|
|
assert out["mask"].shape == (1, 4, 4)
|
|
np.testing.assert_array_equal(out["class_id"], [5])
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# process_transformers_v4_panoptic_segmentation_result
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestProcessTransformersV4PanopticSegmentationResult:
|
|
"""process_transformers_v4_panoptic_segmentation_result decodes PNG masks."""
|
|
|
|
def test_two_segments_produce_two_boolean_masks(self) -> None:
|
|
"""Two segment entries produce two boolean masks with correct coverage."""
|
|
seg_map = np.zeros((4, 4), dtype=np.uint8)
|
|
seg_map[0:2, :] = 1
|
|
seg_map[2:4, :] = 2
|
|
png_bytes = make_panoptic_png(seg_map)
|
|
seg_result = {
|
|
"png_string": png_bytes,
|
|
"segments_info": [
|
|
{"id": 1, "category_id": 10},
|
|
{"id": 2, "category_id": 20},
|
|
],
|
|
}
|
|
|
|
out = process_transformers_v4_panoptic_segmentation_result(seg_result, None)
|
|
|
|
assert out["mask"].shape == (2, 4, 4)
|
|
np.testing.assert_array_equal(out["class_id"], [10, 20])
|
|
# Segment 1 covers top half
|
|
assert out["mask"][0, 0, 0]
|
|
assert not out["mask"][0, 3, 0]
|
|
|
|
def test_with_id2label_sets_class_names(self) -> None:
|
|
"""Providing id2label populates CLASS_NAME_DATA_FIELD in output data."""
|
|
seg_map = np.ones((2, 2), dtype=np.uint8)
|
|
seg_result = {
|
|
"png_string": make_panoptic_png(seg_map),
|
|
"segments_info": [{"id": 1, "category_id": 0}],
|
|
}
|
|
|
|
out = process_transformers_v4_panoptic_segmentation_result(
|
|
seg_result, {0: "background"}
|
|
)
|
|
|
|
np.testing.assert_array_equal(
|
|
out["data"][CLASS_NAME_DATA_FIELD], ["background"]
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# process_transformers_v5_panoptic_segmentation_result
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestProcessTransformersV5PanopticSegmentationResult:
|
|
"""process_transformers_v5_panoptic_segmentation_result uses unique pixel values."""
|
|
|
|
def test_two_unique_ids_produce_two_masks(self) -> None:
|
|
"""Array with two unique values produces two boolean masks."""
|
|
seg_array = np.array([[0, 0, 1, 1], [0, 0, 1, 1]], dtype=np.int64)
|
|
|
|
out = process_transformers_v5_panoptic_segmentation_result(seg_array, None)
|
|
|
|
assert out["mask"].shape[0] == 2
|
|
np.testing.assert_array_equal(out["class_id"], [0, 1])
|
|
|
|
def test_with_id2label_sets_class_names(self) -> None:
|
|
"""id2label maps unique IDs to class name strings in output data."""
|
|
seg_array = np.array([[3, 3], [5, 5]], dtype=np.int64)
|
|
|
|
out = process_transformers_v5_panoptic_segmentation_result(
|
|
seg_array, {3: "tree", 5: "sky"}
|
|
)
|
|
|
|
np.testing.assert_array_equal(
|
|
out["data"][CLASS_NAME_DATA_FIELD], ["tree", "sky"]
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# process_transformers_v5_semantic_or_instance_segmentation_result
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestProcessTransformersV5SemanticOrInstanceSegmentationResult:
|
|
"""process_transformers_v5_semantic_or_instance_segmentation_result."""
|
|
|
|
def test_two_segments_produce_correct_masks_and_scores(self) -> None:
|
|
"""segments_info entries map to masks, scores, and class_ids correctly."""
|
|
seg_arr = np.zeros((4, 4), dtype=np.int64)
|
|
seg_arr[0:2, :] = 1
|
|
seg_arr[2:4, :] = 2
|
|
seg_result = {
|
|
"segmentation": _FakeDetachTensor(seg_arr),
|
|
"segments_info": [
|
|
{"id": 1, "label_id": 0, "score": 0.9},
|
|
{"id": 2, "label_id": 1, "score": 0.7},
|
|
],
|
|
}
|
|
|
|
out = process_transformers_v5_semantic_or_instance_segmentation_result(
|
|
seg_result, None
|
|
)
|
|
|
|
assert out["mask"].shape == (2, 4, 4)
|
|
np.testing.assert_array_equal(out["class_id"], [0, 1])
|
|
np.testing.assert_allclose(out["confidence"], [0.9, 0.7])
|
|
|
|
@pytest.mark.xfail(
|
|
raises=ValueError,
|
|
reason=(
|
|
"empty segments_info produces masks shape (0,) instead of (0,H,W),"
|
|
" causing mask_to_xyxy to crash — source bug, not a test setup issue"
|
|
),
|
|
strict=True,
|
|
)
|
|
def test_empty_segments_info_returns_zero_detections(self) -> None:
|
|
"""Empty segments_info list should yield zero-length arrays (xfail: bug)."""
|
|
seg_result = {
|
|
"segmentation": _FakeDetachTensor(np.zeros((2, 2), dtype=np.int64)),
|
|
"segments_info": [],
|
|
}
|
|
|
|
out = process_transformers_v5_semantic_or_instance_segmentation_result(
|
|
seg_result, None
|
|
)
|
|
|
|
assert len(out["class_id"]) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# process_transformers_v5_segmentation_result (dispatcher)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestProcessTransformersV5SegmentationResult:
|
|
"""process_transformers_v5_segmentation_result dispatches to the right sub-path."""
|
|
|
|
def test_dict_with_segmentation_key_routes_to_semantic_instance_path(
|
|
self,
|
|
) -> None:
|
|
"""Dict input (not Tensor) routes to semantic/instance sub-processor."""
|
|
seg_arr = np.array([[0, 1], [0, 1]], dtype=np.int64)
|
|
seg_result = {
|
|
"segmentation": _FakeDetachTensor(seg_arr),
|
|
"segments_info": [
|
|
{"id": 0, "label_id": 2, "score": 0.95},
|
|
{"id": 1, "label_id": 3, "score": 0.85},
|
|
],
|
|
}
|
|
|
|
out = process_transformers_v5_segmentation_result(seg_result, None)
|
|
|
|
assert len(out["class_id"]) == 2
|
|
|
|
def test_tensor_like_object_routes_to_panoptic_path(self) -> None:
|
|
"""Object whose class is named 'Tensor' routes to panoptic sub-processor."""
|
|
|
|
class Tensor:
|
|
"""Minimal fake torch.Tensor for the panoptic path."""
|
|
|
|
def __init__(self, arr: np.ndarray) -> None:
|
|
self._arr = arr
|
|
|
|
def cpu(self) -> Tensor:
|
|
"""Return self."""
|
|
return self
|
|
|
|
def detach(self) -> Tensor:
|
|
"""Return self."""
|
|
return self
|
|
|
|
def numpy(self) -> np.ndarray:
|
|
"""Return array."""
|
|
return self._arr
|
|
|
|
seg_array = np.array([[0, 1], [0, 1]], dtype=np.int64)
|
|
tensor_result = Tensor(seg_array)
|
|
|
|
out = process_transformers_v5_segmentation_result(tensor_result, None)
|
|
|
|
# Panoptic path: unique IDs [0, 1] → two masks
|
|
assert len(out["class_id"]) == 2
|