"""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 decodes RGB-encoded panoptic IDs.""" def test_extracts_rgb_channels_as_segment_ids(self) -> None: """RGBA PNG: RGB channels 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_decodes_segment_ids_above_255(self) -> None: """RGB panoptic encoding preserves segment IDs beyond one byte.""" seg_map = np.array([[1, 257], [513, 0]], dtype=np.uint32) 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 handles semantic tensors.""" @pytest.mark.parametrize( ("seg_array", "expected_class_ids"), [ pytest.param( np.array([[0, 0, 1, 1], [0, 2, 2, 0]], dtype=np.int64), np.array([0, 1, 2]), id="preserves-class-zero", ), pytest.param( np.zeros((2, 2), dtype=np.int64), np.array([0]), id="single-zero-class", ), ], ) def test_semantic_tensor_preserves_class_zero( self, seg_array: np.ndarray, expected_class_ids: np.ndarray ) -> None: """Bare tensor semantic maps preserve class id zero.""" expected_count = len(expected_class_ids) out = process_transformers_v5_panoptic_segmentation_result(seg_array, None) assert out["mask"].shape == (expected_count, *seg_array.shape) assert out["xyxy"].shape == (expected_count, 4) np.testing.assert_array_equal(out["class_id"], expected_class_ids) 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"] ) def test_with_id2label_preserves_zero_class_name(self) -> None: """id2label maps class id zero when it appears in a tensor map.""" seg_array = np.array([[0, 0], [1, 1]], dtype=np.int64) out = process_transformers_v5_panoptic_segmentation_result( seg_array, {0: "class-zero", 1: "class-one"} ) np.testing.assert_array_equal(out["class_id"], [0, 1]) np.testing.assert_array_equal( out["data"][CLASS_NAME_DATA_FIELD], ["class-zero", "class-one"] ) # --------------------------------------------------------------------------- # 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]) def test_empty_segments_info_returns_zero_detections(self) -> None: """Empty segments_info list yields zero-length detection arrays.""" 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 assert out["xyxy"].shape == (0, 4) assert out["mask"].shape == (0, 2, 2) assert out["confidence"].shape == (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_semantic_tensor_path(self) -> None: """Object whose class is named 'Tensor' routes to semantic tensor path.""" class Tensor: """Minimal fake torch.Tensor for the semantic tensor 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) np.testing.assert_array_equal(out["class_id"], [0, 1])