diff --git a/src/supervision/detection/core.py b/src/supervision/detection/core.py index a131b6c9..f629e96e 100644 --- a/src/supervision/detection/core.py +++ b/src/supervision/detection/core.py @@ -666,10 +666,10 @@ class Detections: or absent. `detections.tracker_id` is `None` when no predictions carry a tracker ID, or when only a subset do (mixed batch) — in that case all tracker IDs are dropped to - preserve alignment with the bounding boxes. Note: mixed - batches containing both RLE/polygon predictions and box-only - predictions may misalign the `mask` array; this is a - pre-existing limitation not addressed by this fix. + preserve alignment with the bounding boxes. + `detections.mask` is `None` when no predictions include mask + data, or when only a subset do (mixed batch) — in that case + all masks are dropped to preserve alignment. Example: ```python diff --git a/src/supervision/detection/utils/internal.py b/src/supervision/detection/utils/internal.py index 467c84b4..ed84b5cc 100644 --- a/src/supervision/detection/utils/internal.py +++ b/src/supervision/detection/utils/internal.py @@ -50,6 +50,23 @@ def extract_ultralytics_masks(yolov8_results: Any) -> npt.NDArray[np.bool_] | No return cast(npt.NDArray[np.bool_], np.asarray(mask_maps, dtype=bool)) +def _all_present_or_none( + values: list[Any], + label: str, + dtype: npt.DTypeLike, +) -> npt.NDArray[Any] | None: + # Identity check (`v is None`) is required when values may contain numpy arrays: + # `None in values` triggers element-wise comparison and raises ValueError. + missing = sum(v is None for v in values) + if 0 < missing < len(values): + logger.warning( + "Partial %s in batch; dropping all to preserve alignment with xyxy.", label + ) + if not values or missing > 0: + return None + return np.array(values, dtype=dtype) + + def process_roboflow_result( roboflow_result: dict[str, Any], ) -> tuple[ @@ -74,10 +91,14 @@ def process_roboflow_result( Returns: A 6-tuple of ``(xyxy, confidence, class_id, masks, tracker_ids, data)`` where each array is aligned with the others. ``masks`` is ``None`` - when no predictions include mask data. ``tracker_ids`` is ``None`` - when no predictions carry a tracker ID, or when only a subset do - (mixed batch) — in that case all tracker IDs are dropped to preserve - alignment with ``xyxy``. + when no predictions include mask data, or when only a subset do + (mixed batch) — in that case all masks are dropped to preserve + alignment with ``xyxy``. A failed RLE decode is treated identically + to a box-only prediction, so one corrupt RLE payload in an otherwise + fully-masked batch also causes all masks to be dropped. + ``tracker_ids`` is ``None`` when no predictions carry a tracker ID, + or when only a subset do (mixed batch) — in that case all tracker + IDs are dropped to preserve alignment with ``xyxy``. Examples: >>> from supervision.detection.utils.internal import process_roboflow_result @@ -100,7 +121,7 @@ def process_roboflow_result( confidence: list[float] = [] class_id: list[int] = [] class_name: list[str] = [] - masks: list[npt.NDArray[np.bool_]] = [] + masks: list[npt.NDArray[np.bool_] | None] = [] tracker_ids: list[int | None] = [] image_width = int(roboflow_result["image"]["width"]) @@ -151,6 +172,7 @@ def process_roboflow_result( class_id.append(prediction["class_id"]) class_name.append(prediction["class"]) confidence.append(prediction["confidence"]) + masks.append(None) tracker_ids.append(prediction.get("tracker_id")) elif len(prediction["points"]) >= 3: polygon = np.array( @@ -180,18 +202,11 @@ def process_roboflow_result( class_name_arr: npt.NDArray[np.str_] = ( np.array(class_name) if len(class_name) > 0 else np.empty(0, dtype=str) ) - masks_arr: npt.NDArray[np.bool_] | None = ( - np.array(masks, dtype=bool) if len(masks) > 0 else None + masks_arr: npt.NDArray[np.bool_] | None = _all_present_or_none( + masks, "mask", dtype=bool ) - if tracker_ids and 0 < tracker_ids.count(None) < len(tracker_ids): - logger.warning( - "Partial tracker_id in batch; dropping all tracker_ids to preserve " - "alignment with xyxy." - ) - tracker_id_arr: npt.NDArray[np.integer] | None = ( - np.array(tracker_ids, dtype=np.int64) - if tracker_ids and None not in tracker_ids - else None + tracker_id_arr: npt.NDArray[np.integer] | None = _all_present_or_none( + tracker_ids, "tracker_id", dtype=np.int64 ) data: _DetectionDataType = {CLASS_NAME_DATA_FIELD: class_name_arr} diff --git a/tests/detection/test_core.py b/tests/detection/test_core.py index d4033337..69b4f76f 100644 --- a/tests/detection/test_core.py +++ b/tests/detection/test_core.py @@ -1013,6 +1013,46 @@ def test_from_inference_partial_tracker_id_does_not_crash() -> None: assert detections["class_name"] is not None +def test_from_inference_partial_mask_does_not_crash() -> None: + """Results where only some predictions carry a mask must not raise.""" + result = { + "image": {"width": 100, "height": 100}, + "predictions": [ + { + "x": 20, + "y": 20, + "width": 20, + "height": 20, + "confidence": 0.9, + "class": "a", + "class_id": 0, + "points": [ + {"x": 10, "y": 10}, + {"x": 30, "y": 10}, + {"x": 30, "y": 30}, + {"x": 10, "y": 30}, + ], + }, + { + "x": 70, + "y": 70, + "width": 20, + "height": 20, + "confidence": 0.8, + "class": "b", + "class_id": 1, + }, + ], + } + + detections = Detections.from_inference(result) + + # all detections are kept; masks are dropped rather than misaligned + assert len(detections) == 2 + assert detections.mask is None + assert detections.xyxy.shape == (2, 4) + + def test_from_inference_empty_class_name_dtype_matches_non_empty() -> None: """Empty and non-empty results should produce string-kind class_name arrays.""" empty_result = {"predictions": [], "image": {"width": 100, "height": 100}} diff --git a/tests/detection/utils/test_internal.py b/tests/detection/utils/test_internal.py index 293ce332..a2be07d3 100644 --- a/tests/detection/utils/test_internal.py +++ b/tests/detection/utils/test_internal.py @@ -436,17 +436,14 @@ TEST_RLE_NONCONTIGUOUS_MASK[0, 3, 2:4] = True np.array([[0.5, 0.5, 2.5, 2.5], [2.0, 2.0, 4.0, 4.0]]), np.array([0.9, 0.8]), np.array([0, 1]), - # NOTE: known misalignment — masks has 1 entry, xyxy has 2. - # Mixed RLE + box-only batches produce mask arrays shorter than - # xyxy; constructing Detections from this result would raise - # ValueError. This is a pre-existing limitation shared by the - # polygon + box-only path. - TEST_RLE_MASK, + # When only some predictions carry a mask, all masks are dropped + # so the result stays aligned with xyxy (mirrors tracker_id). + None, None, {CLASS_NAME_DATA_FIELD: np.array(["person", "car"])}, ), DoesNotRaise(), - ), # mixed RLE + box-only batch — masks misaligned with xyxy (known limitation) + ), # mixed RLE + box-only batch — masks dropped to preserve alignment pytest.param( { "predictions": [