fix(detection): harden model connectors and mask extraction (#2398)
- `from_tensorflow` scaled boxes in place on the array returned by `.numpy()`, which can share memory with the source tensor — corrupting caller data and double-scaling on a repeat call; copy before scaling - `from_lmm` raised a bare `KeyError` for `MOONDREAM` and `QWEN_3_VL`, which the enum and docstring advertise; map both to their `VLM` members - `from_deepseek_vl_2` returned a `(0,)`-shaped `xyxy` on empty output, so a zero-detection response crashed the `Detections` constructor; return `(0, 4)` like the other parsers - `extract_ultralytics_masks` binarized bilinear-resized masks with `> 0`, dilating every mask at object boundaries; threshold at 0.5 to match Ultralytics - add connector coverage: fake-result shims and round-trip tests (N>1, N=1, empty) for the nine previously untested `from_*` connectors and the `detection/tools/transformers.py` processors; one empty-`segments_info` panoptic case is xfail-marked pending a separate fix * test(ci-fix): drop deprecated Pillow mode arg from panoptic helpers * test(coverage): add from_qwen_3_vl end-to-end parametrized tests * test(quality): harden test isolation, xfail strictness, and kwarg forwarding * fix(detection): fix class_name empty dtype; annotate mask threshold asymmetry * chore: ruff-format cleanup (blank lines) --------- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
This commit is contained in:
Jirka Borovec
GitHub
parent
afcf13a6f5
commit
eea04b3656
|
|
@ -453,7 +453,9 @@ class Detections:
|
|||
|
||||
# Tensorflow returns normalized boxes as [ymin, xmin, ymax, xmax], so the
|
||||
# y coordinates (cols 0, 2) scale by height and x (cols 1, 3) by width.
|
||||
boxes = tensorflow_results["detection_boxes"][0].numpy()
|
||||
# `.numpy()` may share memory with the source tensor, so copy before the
|
||||
# in-place scaling to avoid mutating the caller's result / double-scaling.
|
||||
boxes = tensorflow_results["detection_boxes"][0].numpy().copy()
|
||||
boxes[:, [0, 2]] *= resolution_wh[1]
|
||||
boxes[:, [1, 3]] *= resolution_wh[0]
|
||||
boxes = boxes[:, [1, 0, 3, 2]]
|
||||
|
|
@ -1092,6 +1094,7 @@ class Detections:
|
|||
| Google Gemini 2.5 | `GOOGLE_GEMINI_2_5` | detection, segmentation | `resolution_wh` | `classes` |
|
||||
| Moondream | `MOONDREAM` | detection | `resolution_wh` | |
|
||||
| DeepSeek-VL2 | `DEEPSEEK_VL_2` | detection | `resolution_wh` | `classes` |
|
||||
| Qwen3-VL | `QWEN_3_VL` | detection | `resolution_wh` | `classes` |
|
||||
|
||||
Args:
|
||||
lmm: The type of LMM (Large Multimodal Model) to use.
|
||||
|
|
@ -1535,9 +1538,11 @@ class Detections:
|
|||
LMM.PALIGEMMA: VLM.PALIGEMMA,
|
||||
LMM.FLORENCE_2: VLM.FLORENCE_2,
|
||||
LMM.QWEN_2_5_VL: VLM.QWEN_2_5_VL,
|
||||
LMM.QWEN_3_VL: VLM.QWEN_3_VL,
|
||||
LMM.DEEPSEEK_VL_2: VLM.DEEPSEEK_VL_2,
|
||||
LMM.GOOGLE_GEMINI_2_0: VLM.GOOGLE_GEMINI_2_0,
|
||||
LMM.GOOGLE_GEMINI_2_5: VLM.GOOGLE_GEMINI_2_5,
|
||||
LMM.MOONDREAM: VLM.MOONDREAM,
|
||||
}
|
||||
|
||||
if isinstance(lmm, LMM):
|
||||
|
|
@ -1574,7 +1579,7 @@ class Detections:
|
|||
| PaliGemma | `PALIGEMMA` | detection | `resolution_wh` | `classes` |
|
||||
| PaliGemma 2 | `PALIGEMMA` | detection | `resolution_wh` | `classes` |
|
||||
| Qwen2.5-VL | `QWEN_2_5_VL` | detection | `resolution_wh`, `input_wh` | `classes` |
|
||||
| Qwen3-VL | `QWEN_3_VL` | detection | `resolution_wh`, | `classes` |
|
||||
| Qwen3-VL | `QWEN_3_VL` | detection | `resolution_wh` | `classes` |
|
||||
| Google Gemini 2.0 | `GOOGLE_GEMINI_2_0` | detection | `resolution_wh` | `classes` |
|
||||
| Google Gemini 2.5 | `GOOGLE_GEMINI_2_5` | detection, segmentation | `resolution_wh` | `classes` |
|
||||
| Moondream | `MOONDREAM` | detection | `resolution_wh` | |
|
||||
|
|
|
|||
|
|
@ -38,6 +38,20 @@ def _valid_rle_payload(prediction: dict[str, Any]) -> dict[str, Any] | None:
|
|||
|
||||
|
||||
def extract_ultralytics_masks(yolov8_results: Any) -> npt.NDArray[np.bool_] | None:
|
||||
"""Extract boolean masks from Ultralytics results, cropping letterbox padding.
|
||||
|
||||
Handles the case where the inference resolution differs from the original image
|
||||
shape by computing the letterbox padding offsets, cropping them out, and resizing
|
||||
each proto mask back to `orig_shape`. Thresholds at 0.5 to match Ultralytics'
|
||||
semantics and avoid dilating masks through float interpolation.
|
||||
|
||||
Args:
|
||||
yolov8_results: Ultralytics results object with `.masks` and `.orig_shape`.
|
||||
|
||||
Returns:
|
||||
Boolean array of shape `(N, H, W)` aligned with the detections, or `None`
|
||||
when no masks are present.
|
||||
"""
|
||||
if not yolov8_results.masks:
|
||||
return None
|
||||
|
||||
|
|
@ -66,7 +80,12 @@ def extract_ultralytics_masks(yolov8_results: Any) -> npt.NDArray[np.bool_] | No
|
|||
mask = mask[top:bottom, left:right]
|
||||
|
||||
if mask.shape != orig_shape:
|
||||
mask = cv2.resize(mask, (orig_shape[1], orig_shape[0]))
|
||||
# `cv2.resize` interpolates the float proto mask, so threshold at 0.5
|
||||
# (matching Ultralytics' own semantics) instead of casting every
|
||||
# nonzero interpolated value to True, which would dilate the mask.
|
||||
mask = cv2.resize(mask, (orig_shape[1], orig_shape[0])) > 0.5
|
||||
# else: slice-crop (no interpolation) preserves the binary 0/1 float values
|
||||
# produced by Ultralytics; the final np.asarray(..., dtype=bool) is equivalent.
|
||||
|
||||
mask_maps.append(mask)
|
||||
|
||||
|
|
|
|||
|
|
@ -447,7 +447,9 @@ def from_deepseek_vl_2(
|
|||
A tuple of `(xyxy, class_id, class_name)` where `xyxy` is an array of
|
||||
shape `(n, 4)` in format `[x1, y1, x2, y2]`, `class_id` is an
|
||||
optional array of shape `(n,)` with class indices, and `class_name`
|
||||
is an array of shape `(n,)` with class labels.
|
||||
is an array of shape `(n,)` with class labels. When the input
|
||||
contains no detections (or all are filtered by `classes`), returns
|
||||
`(np.empty((0, 4)), np.empty(0), np.empty(0))`.
|
||||
""" # noqa: E501
|
||||
|
||||
width, height = resolution_wh
|
||||
|
|
@ -476,8 +478,14 @@ def from_deepseek_vl_2(
|
|||
)
|
||||
class_name_list.append(current_class_name)
|
||||
|
||||
xyxy = np.array(xyxy_list, dtype=np.float32)
|
||||
class_name = np.array(class_name_list)
|
||||
xyxy = (
|
||||
np.array(xyxy_list, dtype=np.float32)
|
||||
if xyxy_list
|
||||
else np.empty((0, 4), dtype=np.float32)
|
||||
)
|
||||
class_name = (
|
||||
np.array(class_name_list) if class_name_list else np.array([], dtype=object)
|
||||
)
|
||||
|
||||
if classes is not None:
|
||||
mask = np.array([name in classes for name in class_name], dtype=bool)
|
||||
|
|
|
|||
|
|
@ -3,14 +3,23 @@ import pytest
|
|||
|
||||
import supervision.detection.core as detection_core
|
||||
from supervision.config import CLASS_NAME_DATA_FIELD
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.core import LMM, Detections
|
||||
from supervision.detection.vlm import VLM
|
||||
from supervision.utils.internal import SupervisionWarnings
|
||||
from tests.helpers import (
|
||||
_FakeDeepSparseResults,
|
||||
_FakeDetachTensor,
|
||||
_FakeDetectron2Instances,
|
||||
_FakeMMDetPredInstances,
|
||||
_FakeMMDetResults,
|
||||
_FakeNCNNObject,
|
||||
_FakeTensor,
|
||||
_FakeUltralyticsBoxes,
|
||||
_FakeUltralyticsResults,
|
||||
_FakeYoloNasPrediction,
|
||||
_FakeYoloNasResults,
|
||||
_FakeYOLOv5Results,
|
||||
make_panoptic_png,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -128,3 +137,600 @@ def test_from_tensorflow_scales_axes_on_non_square_image() -> None:
|
|||
np.testing.assert_allclose(det.xyxy, [[200.0, 50.0, 600.0, 250.0]])
|
||||
np.testing.assert_allclose(det.confidence, [0.9])
|
||||
np.testing.assert_array_equal(det.class_id, [1])
|
||||
|
||||
|
||||
def test_from_tensorflow_does_not_mutate_source_boxes() -> None:
|
||||
"""Scaling must copy the tensor buffer, leaving the caller's boxes untouched."""
|
||||
source_boxes = np.array([[0.1, 0.2, 0.5, 0.6]], dtype=np.float32)
|
||||
original = source_boxes.copy()
|
||||
results = {
|
||||
"detection_boxes": [_FakeTensor(source_boxes)],
|
||||
"detection_scores": [_FakeTensor(np.array([0.9], dtype=np.float32))],
|
||||
"detection_classes": [_FakeTensor(np.array([1], dtype=np.float32))],
|
||||
}
|
||||
|
||||
det = Detections.from_tensorflow(results, resolution_wh=(1000, 500))
|
||||
|
||||
np.testing.assert_array_equal(source_boxes, original)
|
||||
np.testing.assert_allclose(det.xyxy, [[200.0, 50.0, 600.0, 250.0]])
|
||||
|
||||
|
||||
class TestFromLMMMapping:
|
||||
"""`from_lmm` must map every LMM member to a VLM without raising KeyError."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"lmm_member",
|
||||
[pytest.param(member, id=member.name.lower()) for member in LMM],
|
||||
)
|
||||
def test_from_lmm_maps_every_member_to_vlm(
|
||||
self, lmm_member: LMM, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Each LMM member dispatches to the VLM sharing its value with args intact."""
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_from_vlm(vlm: VLM, result: str, **kwargs: object) -> Detections:
|
||||
captured["vlm"] = vlm
|
||||
captured["result"] = result
|
||||
captured["kwargs"] = kwargs
|
||||
return Detections.empty()
|
||||
|
||||
monkeypatch.setattr(Detections, "from_vlm", staticmethod(fake_from_vlm))
|
||||
|
||||
Detections.from_lmm(lmm_member, result="sentinel", resolution_wh=(10, 10))
|
||||
|
||||
assert isinstance(captured["vlm"], VLM)
|
||||
assert captured["vlm"].value == lmm_member.value # type: ignore[union-attr]
|
||||
assert captured["result"] == "sentinel"
|
||||
assert captured["kwargs"]["resolution_wh"] == (10, 10) # type: ignore[index]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# from_transformers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFromTransformers:
|
||||
"""from_transformers routes detection/segmentation inputs to the right processor."""
|
||||
|
||||
def test_detection_path_maps_boxes_labels_scores(self) -> None:
|
||||
"""Detection result with boxes+labels+scores sets xyxy, class_id, confidence."""
|
||||
xyxy = np.array([[10, 20, 30, 40], [5, 6, 7, 8]], dtype=np.float32)
|
||||
labels = np.array([1, 0], dtype=np.int64)
|
||||
scores = np.array([0.9, 0.5], dtype=np.float32)
|
||||
result = {
|
||||
"boxes": _FakeDetachTensor(xyxy),
|
||||
"labels": _FakeDetachTensor(labels),
|
||||
"scores": _FakeDetachTensor(scores),
|
||||
}
|
||||
|
||||
det = Detections.from_transformers(result)
|
||||
|
||||
np.testing.assert_allclose(det.xyxy, xyxy)
|
||||
np.testing.assert_array_equal(det.class_id, labels.astype(int))
|
||||
np.testing.assert_allclose(det.confidence, scores)
|
||||
|
||||
def test_detection_path_empty_returns_zero_detections(self) -> None:
|
||||
"""Detection path with zero-row tensors yields an empty Detections."""
|
||||
result = {
|
||||
"boxes": _FakeDetachTensor(np.empty((0, 4), dtype=np.float32)),
|
||||
"labels": _FakeDetachTensor(np.empty(0, dtype=np.int64)),
|
||||
"scores": _FakeDetachTensor(np.empty(0, dtype=np.float32)),
|
||||
}
|
||||
|
||||
det = Detections.from_transformers(result)
|
||||
|
||||
assert len(det) == 0
|
||||
|
||||
def test_detection_path_with_id2label_populates_class_names(self) -> None:
|
||||
"""id2label mapping adds class name strings to data dict."""
|
||||
labels = np.array([0, 1], dtype=np.int64)
|
||||
result = {
|
||||
"boxes": _FakeDetachTensor(np.zeros((2, 4), dtype=np.float32)),
|
||||
"labels": _FakeDetachTensor(labels),
|
||||
"scores": _FakeDetachTensor(np.array([0.8, 0.7], dtype=np.float32)),
|
||||
}
|
||||
|
||||
det = Detections.from_transformers(result, id2label={0: "cat", 1: "dog"})
|
||||
|
||||
np.testing.assert_array_equal(det.data[CLASS_NAME_DATA_FIELD], ["cat", "dog"])
|
||||
|
||||
def test_v4_segmentation_masks_without_boxes_yield_correct_shape(self) -> None:
|
||||
"""V4 segmentation with masks-only produces mask of shape (N, H, W)."""
|
||||
masks_bool = np.zeros((2, 4, 4), dtype=bool)
|
||||
masks_bool[0, 0:2, 0:2] = True
|
||||
masks_bool[1, 2:4, 2:4] = True
|
||||
result = {
|
||||
"masks": _FakeDetachTensor(masks_bool.astype(np.uint8)),
|
||||
"labels": _FakeDetachTensor(np.array([0, 1], dtype=np.int64)),
|
||||
"scores": _FakeDetachTensor(np.array([0.9, 0.8], dtype=np.float32)),
|
||||
}
|
||||
|
||||
det = Detections.from_transformers(result)
|
||||
|
||||
assert len(det) == 2
|
||||
assert det.mask is not None
|
||||
assert det.mask.shape == (2, 4, 4)
|
||||
|
||||
def test_v4_panoptic_png_string_extracts_masks_from_red_channel(self) -> None:
|
||||
"""V4 panoptic: png_string+segments_info produce masks keyed by segment id."""
|
||||
seg_map = np.zeros((4, 4), dtype=np.uint8)
|
||||
seg_map[0:2, 0:2] = 1
|
||||
seg_map[2:4, 2:4] = 2
|
||||
png_bytes = make_panoptic_png(seg_map)
|
||||
result = {
|
||||
"png_string": png_bytes,
|
||||
"segments_info": [
|
||||
{"id": 1, "category_id": 3},
|
||||
{"id": 2, "category_id": 7},
|
||||
],
|
||||
}
|
||||
|
||||
det = Detections.from_transformers(result)
|
||||
|
||||
assert len(det) == 2
|
||||
np.testing.assert_array_equal(det.class_id, [3, 7])
|
||||
assert det.mask is not None
|
||||
|
||||
def test_v5_segmentation_key_routes_to_semantic_instance_processor(self) -> None:
|
||||
"""'segmentation' key routes through v5 semantic/instance segmentation path."""
|
||||
seg_arr = np.zeros((4, 4), dtype=np.int64)
|
||||
seg_arr[2:4, :] = 1
|
||||
segments_info = [
|
||||
{"id": 0, "label_id": 0, "score": 0.9},
|
||||
{"id": 1, "label_id": 1, "score": 0.8},
|
||||
]
|
||||
result = {
|
||||
"segmentation": _FakeDetachTensor(seg_arr),
|
||||
"segments_info": segments_info,
|
||||
}
|
||||
|
||||
det = Detections.from_transformers(result)
|
||||
|
||||
assert len(det) == 2
|
||||
np.testing.assert_array_equal(det.class_id, [0, 1])
|
||||
np.testing.assert_allclose(det.confidence, [0.9, 0.8])
|
||||
|
||||
def test_unrecognised_keys_raise_value_error(self) -> None:
|
||||
"""Dict with no valid keys (no boxes/masks/segmentation) raises ValueError."""
|
||||
with pytest.raises(ValueError, match="valid fields"):
|
||||
Detections.from_transformers({})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# from_detectron2
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFromDetectron2:
|
||||
"""from_detectron2 maps Detectron2 pred_instances to Detections fields."""
|
||||
|
||||
def test_two_detections_without_masks_maps_fields(self) -> None:
|
||||
"""N=2, no pred_masks: xyxy, confidence, class_id extracted correctly."""
|
||||
xyxy = np.array([[0, 0, 10, 10], [5, 5, 15, 15]], dtype=np.float32)
|
||||
scores = np.array([0.9, 0.7], dtype=np.float32)
|
||||
class_ids = np.array([1, 2], dtype=np.int64)
|
||||
instances = _FakeDetectron2Instances(xyxy, scores, class_ids)
|
||||
result = {"instances": instances}
|
||||
|
||||
det = Detections.from_detectron2(result)
|
||||
|
||||
np.testing.assert_allclose(det.xyxy, xyxy)
|
||||
np.testing.assert_allclose(det.confidence, scores)
|
||||
np.testing.assert_array_equal(det.class_id, class_ids.astype(int))
|
||||
assert det.mask is None
|
||||
|
||||
def test_single_detection_without_masks(self) -> None:
|
||||
"""N=1 detection without masks returns one-element Detections."""
|
||||
xyxy = np.array([[1, 2, 3, 4]], dtype=np.float32)
|
||||
instances = _FakeDetectron2Instances(
|
||||
xyxy, np.array([0.5], dtype=np.float32), np.array([0])
|
||||
)
|
||||
|
||||
det = Detections.from_detectron2({"instances": instances})
|
||||
|
||||
assert len(det) == 1
|
||||
|
||||
def test_with_pred_masks_sets_mask_field(self) -> None:
|
||||
"""When pred_masks present, mask is populated with boolean array."""
|
||||
xyxy = np.array([[0, 0, 4, 4]], dtype=np.float32)
|
||||
masks = np.ones((1, 4, 4), dtype=bool)
|
||||
instances = _FakeDetectron2Instances(
|
||||
xyxy, np.array([0.8], dtype=np.float32), np.array([0]), masks=masks
|
||||
)
|
||||
|
||||
det = Detections.from_detectron2({"instances": instances})
|
||||
|
||||
assert det.mask is not None
|
||||
assert det.mask.shape == (1, 4, 4)
|
||||
|
||||
def test_empty_instances_returns_zero_length(self) -> None:
|
||||
"""Empty pred_instances arrays produce a zero-length Detections."""
|
||||
instances = _FakeDetectron2Instances(
|
||||
np.empty((0, 4), dtype=np.float32),
|
||||
np.empty(0, dtype=np.float32),
|
||||
np.empty(0, dtype=np.int64),
|
||||
)
|
||||
|
||||
det = Detections.from_detectron2({"instances": instances})
|
||||
|
||||
assert len(det) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# from_mmdetection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFromMMDetection:
|
||||
"""from_mmdetection maps MMDet pred_instances to Detections fields."""
|
||||
|
||||
def test_two_detections_without_masks(self) -> None:
|
||||
"""N=2, no masks attribute: xyxy, confidence, class_id set; mask is None."""
|
||||
xyxy = np.array([[0, 0, 10, 10], [5, 5, 20, 20]], dtype=np.float32)
|
||||
scores = np.array([0.85, 0.6], dtype=np.float32)
|
||||
labels = np.array([0, 3], dtype=np.int64)
|
||||
pred_instances = _FakeMMDetPredInstances(xyxy, scores, labels)
|
||||
result = _FakeMMDetResults(pred_instances)
|
||||
|
||||
det = Detections.from_mmdetection(result)
|
||||
|
||||
np.testing.assert_allclose(det.xyxy, xyxy)
|
||||
np.testing.assert_allclose(det.confidence, scores)
|
||||
np.testing.assert_array_equal(det.class_id, labels.astype(int))
|
||||
assert det.mask is None
|
||||
|
||||
def test_single_detection_without_masks(self) -> None:
|
||||
"""N=1 detection without masks returns one-element Detections."""
|
||||
pred_instances = _FakeMMDetPredInstances(
|
||||
np.array([[2, 4, 6, 8]], dtype=np.float32),
|
||||
np.array([0.75], dtype=np.float32),
|
||||
np.array([1]),
|
||||
)
|
||||
|
||||
det = Detections.from_mmdetection(_FakeMMDetResults(pred_instances))
|
||||
|
||||
assert len(det) == 1
|
||||
|
||||
def test_with_masks_populates_mask_field(self) -> None:
|
||||
"""When masks present in pred_instances, mask field is set."""
|
||||
xyxy = np.array([[0, 0, 4, 4]], dtype=np.float32)
|
||||
masks = np.ones((1, 4, 4), dtype=bool)
|
||||
pred_instances = _FakeMMDetPredInstances(
|
||||
xyxy, np.array([0.9], dtype=np.float32), np.array([0]), masks=masks
|
||||
)
|
||||
|
||||
det = Detections.from_mmdetection(_FakeMMDetResults(pred_instances))
|
||||
|
||||
assert det.mask is not None
|
||||
assert det.mask.shape == (1, 4, 4)
|
||||
|
||||
def test_empty_pred_instances_returns_zero_length(self) -> None:
|
||||
"""Empty bboxes/scores/labels arrays yield zero-length Detections."""
|
||||
pred_instances = _FakeMMDetPredInstances(
|
||||
np.empty((0, 4), dtype=np.float32),
|
||||
np.empty(0, dtype=np.float32),
|
||||
np.empty(0, dtype=np.int64),
|
||||
)
|
||||
|
||||
det = Detections.from_mmdetection(_FakeMMDetResults(pred_instances))
|
||||
|
||||
assert len(det) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# from_paddledet
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFromPaddleDet:
|
||||
"""from_paddledet extracts xyxy, confidence, class_id from the bbox column array."""
|
||||
|
||||
def test_empty_bbox_returns_empty_detections(self) -> None:
|
||||
"""Empty (0,6) bbox array yields zero-length Detections."""
|
||||
det = Detections.from_paddledet({"bbox": np.empty((0, 6), dtype=np.float32)})
|
||||
|
||||
assert len(det) == 0
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("bbox_array", "expected_len"),
|
||||
[
|
||||
pytest.param(
|
||||
np.array(
|
||||
[[0, 0.9, 10, 20, 30, 40], [1, 0.7, 5, 6, 7, 8]],
|
||||
dtype=np.float32,
|
||||
),
|
||||
2,
|
||||
id="two-detections",
|
||||
),
|
||||
pytest.param(
|
||||
np.array([[2, 0.5, 1, 2, 3, 4]], dtype=np.float32),
|
||||
1,
|
||||
id="single-detection",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_maps_bbox_columns_to_detections(
|
||||
self, bbox_array: np.ndarray, expected_len: int
|
||||
) -> None:
|
||||
"""bbox[:,0]=class_id, [:,1]=confidence, [:,2:6]=xyxy are extracted."""
|
||||
result = {"bbox": bbox_array}
|
||||
|
||||
det = Detections.from_paddledet(result)
|
||||
|
||||
assert len(det) == expected_len
|
||||
np.testing.assert_allclose(det.xyxy, bbox_array[:, 2:6])
|
||||
np.testing.assert_allclose(det.confidence, bbox_array[:, 1])
|
||||
np.testing.assert_array_equal(det.class_id, bbox_array[:, 0].astype(int))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# from_deepsparse
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFromDeepSparse:
|
||||
"""from_deepsparse maps DeepSparse boxes/scores/labels to Detections."""
|
||||
|
||||
def test_empty_results_return_empty_detections(self) -> None:
|
||||
"""Empty boxes/scores/labels arrays yield zero-length Detections."""
|
||||
result = _FakeDeepSparseResults(
|
||||
boxes=[np.empty((0, 4), dtype=np.float32)],
|
||||
scores=[np.empty(0, dtype=np.float32)],
|
||||
labels=[np.empty(0, dtype=np.float32)],
|
||||
)
|
||||
|
||||
det = Detections.from_deepsparse(result)
|
||||
|
||||
assert len(det) == 0
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("boxes", "scores", "labels", "expected_len"),
|
||||
[
|
||||
pytest.param(
|
||||
np.array([[0, 0, 10, 10], [5, 5, 15, 15]], dtype=np.float32),
|
||||
np.array([0.95, 0.8], dtype=np.float32),
|
||||
np.array([0, 1], dtype=np.float32),
|
||||
2,
|
||||
id="two-detections",
|
||||
),
|
||||
pytest.param(
|
||||
np.array([[1, 2, 3, 4]], dtype=np.float32),
|
||||
np.array([0.6], dtype=np.float32),
|
||||
np.array([3], dtype=np.float32),
|
||||
1,
|
||||
id="single-detection",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_maps_boxes_scores_labels_to_detections(
|
||||
self,
|
||||
boxes: np.ndarray,
|
||||
scores: np.ndarray,
|
||||
labels: np.ndarray,
|
||||
expected_len: int,
|
||||
) -> None:
|
||||
"""boxes[0], scores[0], labels[0] are extracted into Detections fields."""
|
||||
result = _FakeDeepSparseResults(boxes=[boxes], scores=[scores], labels=[labels])
|
||||
|
||||
det = Detections.from_deepsparse(result)
|
||||
|
||||
assert len(det) == expected_len
|
||||
np.testing.assert_allclose(det.xyxy, boxes)
|
||||
np.testing.assert_allclose(det.confidence, scores)
|
||||
np.testing.assert_array_equal(det.class_id, labels.astype(int))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# from_easyocr
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFromEasyOCR:
|
||||
"""from_easyocr converts EasyOCR polygon-corner results to Detections."""
|
||||
|
||||
def test_two_detections_produces_correct_xyxy_and_text(self) -> None:
|
||||
"""Two detections with rectangular corners produce correct xyxy and text."""
|
||||
bbox_a = [[10, 10], [30, 10], [30, 20], [10, 20]]
|
||||
bbox_b = [[50, 5], [80, 5], [80, 25], [50, 25]]
|
||||
results = [
|
||||
(bbox_a, "hello", 0.95),
|
||||
(bbox_b, "world", 0.80),
|
||||
]
|
||||
|
||||
det = Detections.from_easyocr(results)
|
||||
|
||||
assert len(det) == 2
|
||||
np.testing.assert_allclose(det.xyxy[0], [10, 10, 30, 20])
|
||||
np.testing.assert_allclose(det.xyxy[1], [50, 5, 80, 25])
|
||||
np.testing.assert_array_equal(
|
||||
det.data[CLASS_NAME_DATA_FIELD], ["hello", "world"]
|
||||
)
|
||||
|
||||
def test_single_detection_returns_one_element(self) -> None:
|
||||
"""N=1 result returns a single-detection Detections."""
|
||||
results = [([[0, 0], [10, 0], [10, 5], [0, 5]], "ok", 0.7)]
|
||||
|
||||
det = Detections.from_easyocr(results)
|
||||
|
||||
assert len(det) == 1
|
||||
|
||||
def test_empty_list_returns_empty_detections(self) -> None:
|
||||
"""Empty input returns an empty Detections."""
|
||||
det = Detections.from_easyocr([])
|
||||
|
||||
assert len(det) == 0
|
||||
|
||||
def test_missing_confidence_defaults_to_zero(self) -> None:
|
||||
"""Two-element tuples (no confidence) default confidence to 0."""
|
||||
results = [([[0, 0], [5, 0], [5, 3], [0, 3]], "hi")]
|
||||
|
||||
det = Detections.from_easyocr(results)
|
||||
|
||||
assert len(det) == 1
|
||||
assert float(det.confidence[0]) == pytest.approx(0.0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# from_azure_analyze_image
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_azure_result(
|
||||
detections: list[dict],
|
||||
) -> dict:
|
||||
"""Build a minimal Azure Image Analysis response dict."""
|
||||
return {"objectsResult": {"values": detections}}
|
||||
|
||||
|
||||
def _make_azure_detection(x: int, y: int, w: int, h: int, tags: list[dict]) -> dict:
|
||||
"""Build one Azure detection entry."""
|
||||
return {
|
||||
"boundingBox": {"x": x, "y": y, "w": w, "h": h},
|
||||
"tags": tags,
|
||||
}
|
||||
|
||||
|
||||
class TestFromAzureAnalyzeImage:
|
||||
"""from_azure_analyze_image converts Azure object detection results."""
|
||||
|
||||
def test_dynamic_class_mapping_assigns_ids_in_order(self) -> None:
|
||||
"""Without class_map, unique class names get monotonically increasing IDs."""
|
||||
result = _make_azure_result(
|
||||
[
|
||||
_make_azure_detection(
|
||||
0, 0, 10, 10, [{"name": "cat", "confidence": 0.9}]
|
||||
),
|
||||
_make_azure_detection(
|
||||
20, 20, 10, 10, [{"name": "dog", "confidence": 0.7}]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
det = Detections.from_azure_analyze_image(result)
|
||||
|
||||
assert len(det) == 2
|
||||
# cat gets id=0 (first seen), dog gets id=1
|
||||
np.testing.assert_array_equal(det.class_id, [0, 1])
|
||||
np.testing.assert_allclose(det.confidence, [0.9, 0.7])
|
||||
np.testing.assert_allclose(det.xyxy[0], [0, 0, 10, 10])
|
||||
|
||||
def test_explicit_class_map_filters_unknown_classes(self) -> None:
|
||||
"""With class_map, tags whose name is absent from the map are dropped."""
|
||||
class_map = {5: "cat"}
|
||||
result = _make_azure_result(
|
||||
[
|
||||
_make_azure_detection(
|
||||
0,
|
||||
0,
|
||||
10,
|
||||
10,
|
||||
[
|
||||
{"name": "cat", "confidence": 0.9},
|
||||
{"name": "unknown", "confidence": 0.5},
|
||||
],
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
det = Detections.from_azure_analyze_image(result, class_map=class_map)
|
||||
|
||||
# Only 'cat' (id=5) survives; 'unknown' is filtered
|
||||
assert len(det) == 1
|
||||
assert int(det.class_id[0]) == 5
|
||||
|
||||
def test_empty_values_list_returns_empty_detections(self) -> None:
|
||||
"""Zero detections in values list produce an empty Detections."""
|
||||
result = _make_azure_result([])
|
||||
|
||||
det = Detections.from_azure_analyze_image(result)
|
||||
|
||||
assert len(det) == 0
|
||||
|
||||
def test_error_key_raises_value_error(self) -> None:
|
||||
"""Response containing 'error' key raises ValueError."""
|
||||
result = {"error": {"message": "service unavailable"}}
|
||||
|
||||
with pytest.raises(ValueError, match="service unavailable"):
|
||||
Detections.from_azure_analyze_image(result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# from_ncnn
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFromNCNN:
|
||||
"""from_ncnn converts ncnn rect objects (xywh) to xyxy Detections."""
|
||||
|
||||
def test_empty_objects_return_empty_detections(self) -> None:
|
||||
"""Empty object list yields zero-length Detections."""
|
||||
det = Detections.from_ncnn([])
|
||||
|
||||
assert len(det) == 0
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("objects", "expected_len"),
|
||||
[
|
||||
pytest.param(
|
||||
[
|
||||
_FakeNCNNObject(10, 20, 30, 40, 0.9, 0),
|
||||
_FakeNCNNObject(5, 5, 10, 10, 0.7, 1),
|
||||
],
|
||||
2,
|
||||
id="two-detections",
|
||||
),
|
||||
pytest.param(
|
||||
[_FakeNCNNObject(0, 0, 20, 20, 0.5, 2)],
|
||||
1,
|
||||
id="single-detection",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_maps_xywh_rect_to_xyxy(self, objects: list, expected_len: int) -> None:
|
||||
"""rect xywh converts to xyxy; prob and label map to confidence/class_id."""
|
||||
det = Detections.from_ncnn(objects)
|
||||
|
||||
assert len(det) == expected_len
|
||||
first = objects[0]
|
||||
expected_x2 = first.rect.x + first.rect.w
|
||||
expected_y2 = first.rect.y + first.rect.h
|
||||
np.testing.assert_allclose(det.xyxy[0, 2], expected_x2)
|
||||
np.testing.assert_allclose(det.xyxy[0, 3], expected_y2)
|
||||
assert float(det.confidence[0]) == pytest.approx(first.prob)
|
||||
assert int(det.class_id[0]) == first.label
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# from_lmm end-to-end
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFromLMMEndToEnd:
|
||||
"""from_lmm end-to-end: deprecated dispatcher produces correct Detections."""
|
||||
|
||||
def test_paligemma_result_produces_correct_xyxy(self) -> None:
|
||||
"""PaliGemma loc-token string is correctly parsed through the legacy API."""
|
||||
result = "<loc0256><loc0256><loc0768><loc0768> cat"
|
||||
|
||||
with pytest.warns(SupervisionWarnings):
|
||||
det = Detections.from_lmm(
|
||||
LMM.PALIGEMMA,
|
||||
result,
|
||||
resolution_wh=(1000, 1000),
|
||||
classes=["cat"],
|
||||
)
|
||||
|
||||
assert len(det) == 1
|
||||
np.testing.assert_allclose(det.xyxy, [[250.0, 250.0, 750.0, 750.0]])
|
||||
assert int(det.class_id[0]) == 0
|
||||
|
||||
def test_string_lmm_name_is_accepted_and_dispatches(self) -> None:
|
||||
"""Passing LMM name as lowercase string works identically to the enum."""
|
||||
with pytest.warns(SupervisionWarnings):
|
||||
det = Detections.from_lmm(
|
||||
"paligemma",
|
||||
"",
|
||||
resolution_wh=(1000, 1000),
|
||||
)
|
||||
|
||||
assert len(det) == 0
|
||||
|
|
|
|||
|
|
@ -14,9 +14,59 @@ from supervision.detection.vlm import (
|
|||
from_moondream,
|
||||
from_paligemma,
|
||||
from_qwen_2_5_vl,
|
||||
from_qwen_3_vl,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("result", "resolution_wh", "classes", "expected_xyxy", "expected_class_name"),
|
||||
[
|
||||
pytest.param(
|
||||
'```json\n[{"bbox_2d": [100, 200, 300, 400], "label": "cat"}]\n```',
|
||||
(640, 480),
|
||||
None,
|
||||
np.array([[64.0, 96.0, 192.0, 192.0]]),
|
||||
np.array(["cat"], dtype=str),
|
||||
id="single-detection-scales-from-1000x1000",
|
||||
),
|
||||
pytest.param(
|
||||
"```json\n[]\n```",
|
||||
(640, 480),
|
||||
None,
|
||||
np.empty((0, 4)),
|
||||
np.empty(0, dtype=str),
|
||||
id="empty-json-array-returns-empty",
|
||||
),
|
||||
pytest.param(
|
||||
"```json\n"
|
||||
'[{"bbox_2d": [0, 0, 500, 500], "label": "dog"},'
|
||||
' {"bbox_2d": [500, 500, 1000, 1000], "label": "cat"}]\n```',
|
||||
(640, 480),
|
||||
["cat"],
|
||||
np.array([[320.0, 240.0, 640.0, 480.0]]),
|
||||
np.array(["cat"], dtype=str),
|
||||
id="classes-filter-keeps-only-matching",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_from_qwen_3_vl(
|
||||
result: str,
|
||||
resolution_wh: tuple[int, int],
|
||||
classes: list[str] | None,
|
||||
expected_xyxy: np.ndarray,
|
||||
expected_class_name: np.ndarray,
|
||||
) -> None:
|
||||
"""from_qwen_3_vl scales from implicit 1000x1000 input space to resolution_wh."""
|
||||
xyxy, _class_id, class_name = from_qwen_3_vl(
|
||||
result=result,
|
||||
resolution_wh=resolution_wh,
|
||||
classes=classes,
|
||||
)
|
||||
|
||||
np.testing.assert_allclose(xyxy, expected_xyxy)
|
||||
np.testing.assert_array_equal(class_name, expected_class_name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("exception", "result", "resolution_wh", "classes", "expected_results"),
|
||||
[
|
||||
|
|
@ -1194,19 +1244,19 @@ def test_from_google_gemini_2_5(
|
|||
("exception", "result", "resolution_wh", "classes", "expected_detections"),
|
||||
[
|
||||
(
|
||||
pytest.raises(ValueError, match=r"xyxy must be a 2D np\.ndarray"),
|
||||
does_not_raise(),
|
||||
"",
|
||||
(100, 100),
|
||||
None,
|
||||
None,
|
||||
), # empty text
|
||||
Detections.empty(),
|
||||
), # empty text -> empty detections (aligned with other VLM parsers)
|
||||
(
|
||||
pytest.raises(ValueError, match=r"xyxy must be a 2D np\.ndarray"),
|
||||
does_not_raise(),
|
||||
"random text",
|
||||
(100, 100),
|
||||
None,
|
||||
None,
|
||||
), # random text
|
||||
Detections.empty(),
|
||||
), # random text -> empty detections
|
||||
(
|
||||
does_not_raise(),
|
||||
"<|ref|>cat<|/ref|><|det|>[[100, 200, 300, 400]]<|/det|>",
|
||||
|
|
@ -1297,6 +1347,29 @@ def test_from_deepseek_vl_2(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("result", "classes"),
|
||||
[
|
||||
pytest.param("", None, id="empty_string"),
|
||||
pytest.param("no tags here", None, id="no_tags"),
|
||||
pytest.param("", ["cat"], id="empty_string_with_classes"),
|
||||
],
|
||||
)
|
||||
def test_from_deepseek_vl_2_empty_parse_returns_empty_detections(
|
||||
result: str, classes: list[str] | None
|
||||
) -> None:
|
||||
"""A result with no ref/det pairs yields empty Detections instead of raising."""
|
||||
detections = Detections.from_vlm(
|
||||
vlm=VLM.DEEPSEEK_VL_2,
|
||||
result=result,
|
||||
resolution_wh=(1000, 1000),
|
||||
classes=classes,
|
||||
)
|
||||
|
||||
assert len(detections) == 0
|
||||
assert detections.xyxy.shape == (0, 4)
|
||||
|
||||
|
||||
def test_from_google_gemini_2_5_malformed_mask_keeps_confidence_aligned():
|
||||
"""A non-data-URI mask must not skip the item's confidence and desync arrays."""
|
||||
result = (
|
||||
|
|
|
|||
|
|
@ -0,0 +1,373 @@
|
|||
"""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
|
||||
|
|
@ -1,12 +1,14 @@
|
|||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.config import CLASS_NAME_DATA_FIELD
|
||||
from supervision.detection.compact_mask import CompactMask
|
||||
from supervision.detection.utils.internal import (
|
||||
extract_ultralytics_masks,
|
||||
get_data_item,
|
||||
merge_data,
|
||||
merge_metadata,
|
||||
|
|
@ -14,6 +16,54 @@ from supervision.detection.utils.internal import (
|
|||
)
|
||||
|
||||
|
||||
class _FakeMasksData:
|
||||
"""Ultralytics-like mask tensor exposing shape and cpu().numpy()."""
|
||||
|
||||
def __init__(self, arr: np.ndarray) -> None:
|
||||
self._arr = np.asarray(arr, dtype=np.float32)
|
||||
self.shape = self._arr.shape
|
||||
|
||||
def cpu(self) -> "_FakeMasksData":
|
||||
return self
|
||||
|
||||
def numpy(self) -> np.ndarray:
|
||||
return self._arr
|
||||
|
||||
|
||||
class _FakeMasks:
|
||||
"""Ultralytics-like masks container holding a data tensor."""
|
||||
|
||||
def __init__(self, data: _FakeMasksData) -> None:
|
||||
self.data = data
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class _FakeYOLOMaskResults:
|
||||
"""Minimal Ultralytics results exposing masks and orig_shape."""
|
||||
|
||||
def __init__(self, masks_arr: np.ndarray, orig_shape: tuple[int, int]) -> None:
|
||||
self.masks = _FakeMasks(_FakeMasksData(masks_arr))
|
||||
self.orig_shape = orig_shape
|
||||
|
||||
|
||||
def test_extract_ultralytics_masks_thresholds_resized_proto_at_half() -> None:
|
||||
"""Resized proto masks threshold at 0.5, so interpolated edges don't dilate."""
|
||||
orig_shape = (4, 4)
|
||||
proto = np.array([[[1.0, 1.0], [0.0, 0.0]]], dtype=np.float32)
|
||||
results = _FakeYOLOMaskResults(masks_arr=proto, orig_shape=orig_shape)
|
||||
|
||||
masks = extract_ultralytics_masks(results)
|
||||
|
||||
resized = cv2.resize(proto[0], (orig_shape[1], orig_shape[0]))
|
||||
assert masks is not None
|
||||
assert masks.dtype == bool
|
||||
np.testing.assert_array_equal(masks[0], resized > 0.5)
|
||||
# A naive `> 0` cast would flag the interpolated boundary row as True.
|
||||
assert masks[0].sum() < int((resized > 0).sum())
|
||||
|
||||
|
||||
def _pred(
|
||||
yx: tuple[float, float] = (1.5, 1.5),
|
||||
size: tuple[float, float] = (2.0, 2.0),
|
||||
|
|
|
|||
122
tests/helpers.py
122
tests/helpers.py
|
|
@ -8,14 +8,29 @@ for generating synthetic test data and performing custom assertions.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.key_points.core import KeyPoints
|
||||
|
||||
|
||||
def make_panoptic_png(segment_map: np.ndarray) -> bytes:
|
||||
"""Encode a (H, W) uint8 segment-ID array as a 4-channel RGBA PNG byte string.
|
||||
|
||||
The segment IDs are stored in the red channel (channel 0). Used by
|
||||
panoptic segmentation tests that construct PNG-encoded segment maps.
|
||||
"""
|
||||
arr = np.zeros((*segment_map.shape, 4), dtype=np.uint8)
|
||||
arr[:, :, 0] = segment_map.astype(np.uint8)
|
||||
buf = io.BytesIO()
|
||||
Image.fromarray(arr).save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _create_detections(
|
||||
xyxy: list[list[float]],
|
||||
mask: list[np.ndarray] | None = None,
|
||||
|
|
@ -468,6 +483,113 @@ def create_yolo_dataset(
|
|||
}
|
||||
|
||||
|
||||
class _FakeDetachTensor:
|
||||
"""Fake torch.Tensor supporting the cpu().detach().numpy() call chain."""
|
||||
|
||||
def __init__(self, arr: np.ndarray) -> None:
|
||||
self._arr = np.asarray(arr)
|
||||
|
||||
def cpu(self) -> _FakeDetachTensor:
|
||||
"""Return self to allow chaining."""
|
||||
return self
|
||||
|
||||
def detach(self) -> _FakeDetachTensor:
|
||||
"""Return self to allow chaining."""
|
||||
return self
|
||||
|
||||
def numpy(self) -> np.ndarray:
|
||||
"""Return underlying array."""
|
||||
return self._arr
|
||||
|
||||
|
||||
class _FakeDetectron2Boxes:
|
||||
"""Fake Detectron2 Boxes exposing .tensor for the cpu().numpy() chain."""
|
||||
|
||||
def __init__(self, xyxy: np.ndarray) -> None:
|
||||
self.tensor = _FakeTensor(xyxy)
|
||||
|
||||
|
||||
class _FakeDetectron2Instances:
|
||||
"""Fake Detectron2 Instances: pred_boxes, scores, pred_classes, optional masks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
xyxy: np.ndarray,
|
||||
scores: np.ndarray,
|
||||
class_ids: np.ndarray,
|
||||
masks: np.ndarray | None = None,
|
||||
) -> None:
|
||||
self.pred_boxes = _FakeDetectron2Boxes(xyxy)
|
||||
self.scores = _FakeTensor(scores)
|
||||
self.pred_classes = _FakeTensor(class_ids)
|
||||
if masks is not None:
|
||||
self.pred_masks = _FakeTensor(masks)
|
||||
|
||||
|
||||
class _FakeMMDetPredInstances:
|
||||
"""Fake MMDetection pred_instances supporting the 'masks' in membership check."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
xyxy: np.ndarray,
|
||||
scores: np.ndarray,
|
||||
labels: np.ndarray,
|
||||
masks: np.ndarray | None = None,
|
||||
) -> None:
|
||||
self.bboxes = _FakeTensor(xyxy)
|
||||
self.scores = _FakeTensor(scores)
|
||||
self.labels = _FakeTensor(labels)
|
||||
self._masks: np.ndarray | None = masks
|
||||
if masks is not None:
|
||||
self.masks = _FakeTensor(masks)
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
"""Return True for 'masks' only when masks were provided at construction."""
|
||||
return key == "masks" and self._masks is not None
|
||||
|
||||
|
||||
class _FakeMMDetResults:
|
||||
"""Fake MMDetection inference result wrapping pred_instances."""
|
||||
|
||||
def __init__(self, pred_instances: _FakeMMDetPredInstances) -> None:
|
||||
self.pred_instances = pred_instances
|
||||
|
||||
|
||||
class _FakeDeepSparseResults:
|
||||
"""Fake DeepSparse inference result with list attributes boxes, scores, labels."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
boxes: list[np.ndarray],
|
||||
scores: list[np.ndarray],
|
||||
labels: list[np.ndarray],
|
||||
) -> None:
|
||||
self.boxes = boxes
|
||||
self.scores = scores
|
||||
self.labels = labels
|
||||
|
||||
|
||||
class _FakeNCNNRect:
|
||||
"""Fake ncnn Rect with x, y, w, h as numpy float32 scalars supporting .astype()."""
|
||||
|
||||
def __init__(self, x: float, y: float, w: float, h: float) -> None:
|
||||
self.x = np.float32(x)
|
||||
self.y = np.float32(y)
|
||||
self.w = np.float32(w)
|
||||
self.h = np.float32(h)
|
||||
|
||||
|
||||
class _FakeNCNNObject:
|
||||
"""Fake ncnn detected object with rect, prob, label."""
|
||||
|
||||
def __init__(
|
||||
self, x: float, y: float, w: float, h: float, prob: float, label: int
|
||||
) -> None:
|
||||
self.rect = _FakeNCNNRect(x, y, w, h)
|
||||
self.prob = prob
|
||||
self.label = label
|
||||
|
||||
|
||||
def create_predictions_with_class_iou_tests(
|
||||
gt_detections: Detections, num_classes: int
|
||||
) -> Detections:
|
||||
|
|
|
|||
Loading…
Reference in New Issue