feat: Add soft Non-Max suppression (#1624)
* feat: Add soft Non-Max suppression * feat(nms): add vectorized Gaussian Soft-NMS box/mask primitives * feat(core): add Detections.with_soft_nms * feat: export soft-NMS functions from top-level supervision API * docs: add soft-NMS entries to IoU/NMS utils page * test: add Soft-NMS coverage for box/mask primitives and Detections API --------- Co-authored-by: jirka <6035284+Borda@users.noreply.github.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
This commit is contained in:
HALLOUARD
GitHub
parent
140c05f273
commit
89d49c2e93
|
|
@ -52,12 +52,24 @@ comments: true
|
|||
|
||||
:::supervision.detection.utils.iou_and_nms.box_non_max_suppression
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.iou_and_nms.box_soft_non_max_suppression">box_soft_non_max_suppression</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.iou_and_nms.box_soft_non_max_suppression
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.iou_and_nms.mask_non_max_suppression">mask_non_max_suppression</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.iou_and_nms.mask_non_max_suppression
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.iou_and_nms.mask_soft_non_max_suppression">mask_soft_non_max_suppression</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.iou_and_nms.mask_soft_non_max_suppression
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.iou_and_nms.box_non_max_merge">box_non_max_merge</a></h2>
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -91,9 +91,11 @@ from supervision.detection.utils.iou_and_nms import (
|
|||
box_iou_batch_with_jaccard,
|
||||
box_non_max_merge,
|
||||
box_non_max_suppression,
|
||||
box_soft_non_max_suppression,
|
||||
mask_iou_batch,
|
||||
mask_non_max_merge,
|
||||
mask_non_max_suppression,
|
||||
mask_soft_non_max_suppression,
|
||||
oriented_box_iou_batch,
|
||||
oriented_box_non_max_merge,
|
||||
oriented_box_non_max_suppression,
|
||||
|
|
@ -235,6 +237,7 @@ __all__ = [
|
|||
"box_iou_batch_with_jaccard",
|
||||
"box_non_max_merge",
|
||||
"box_non_max_suppression",
|
||||
"box_soft_non_max_suppression",
|
||||
"calculate_masks_centroids",
|
||||
"calculate_optimal_line_thickness",
|
||||
"calculate_optimal_text_scale",
|
||||
|
|
@ -268,6 +271,7 @@ __all__ = [
|
|||
"mask_iou_batch",
|
||||
"mask_non_max_merge",
|
||||
"mask_non_max_suppression",
|
||||
"mask_soft_non_max_suppression",
|
||||
"mask_to_polygons",
|
||||
"mask_to_rle",
|
||||
"mask_to_roi",
|
||||
|
|
|
|||
|
|
@ -50,9 +50,11 @@ from supervision.detection.utils.iou_and_nms import (
|
|||
box_iou_batch,
|
||||
box_non_max_merge,
|
||||
box_non_max_suppression,
|
||||
box_soft_non_max_suppression,
|
||||
mask_iou_batch,
|
||||
mask_non_max_merge,
|
||||
mask_non_max_suppression,
|
||||
mask_soft_non_max_suppression,
|
||||
oriented_box_non_max_merge,
|
||||
oriented_box_non_max_suppression,
|
||||
)
|
||||
|
|
@ -2971,6 +2973,37 @@ class Detections:
|
|||
)
|
||||
return new
|
||||
|
||||
def _build_nms_predictions(
|
||||
self, class_agnostic: bool, operation_name: str
|
||||
) -> npt.NDArray[np.floating]:
|
||||
"""Stack xyxy + confidence (+ class_id) for NMS/NMM/Soft-NMS dispatch.
|
||||
|
||||
Callers must already have verified `self.confidence is not None`.
|
||||
"""
|
||||
if class_agnostic:
|
||||
return cast(
|
||||
npt.NDArray[np.floating],
|
||||
np.hstack(
|
||||
(self.xyxy, cast(np.ndarray, self.confidence).reshape(-1, 1))
|
||||
),
|
||||
)
|
||||
if self.class_id is None:
|
||||
raise ValueError(
|
||||
f"Detections class_id must be given for {operation_name} to be "
|
||||
f"executed. If you intended to perform class agnostic "
|
||||
f"{operation_name} set class_agnostic=True."
|
||||
)
|
||||
return cast(
|
||||
npt.NDArray[np.floating],
|
||||
np.hstack(
|
||||
(
|
||||
self.xyxy,
|
||||
cast(np.ndarray, self.confidence).reshape(-1, 1),
|
||||
self.class_id.reshape(-1, 1),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
def with_nms(
|
||||
self,
|
||||
threshold: float = 0.5,
|
||||
|
|
@ -3009,28 +3042,7 @@ class Detections:
|
|||
"Detections confidence must be given for NMS to be executed."
|
||||
)
|
||||
|
||||
if class_agnostic:
|
||||
predictions = cast(
|
||||
npt.NDArray[np.floating],
|
||||
np.hstack((self.xyxy, self.confidence.reshape(-1, 1))),
|
||||
)
|
||||
else:
|
||||
if self.class_id is None:
|
||||
raise ValueError(
|
||||
"Detections class_id must be given for NMS to be executed. If "
|
||||
"you intended to perform class agnostic NMS "
|
||||
"set class_agnostic=True."
|
||||
)
|
||||
predictions = cast(
|
||||
npt.NDArray[np.floating],
|
||||
np.hstack(
|
||||
(
|
||||
self.xyxy,
|
||||
self.confidence.reshape(-1, 1),
|
||||
self.class_id.reshape(-1, 1),
|
||||
)
|
||||
),
|
||||
)
|
||||
predictions = self._build_nms_predictions(class_agnostic, "NMS")
|
||||
|
||||
if self.mask is not None:
|
||||
indices = mask_non_max_suppression(
|
||||
|
|
@ -3057,6 +3069,79 @@ class Detections:
|
|||
|
||||
return self.select(indices)
|
||||
|
||||
def with_soft_nms(
|
||||
self,
|
||||
sigma: float = 0.5,
|
||||
class_agnostic: bool = False,
|
||||
score_threshold: float | None = None,
|
||||
) -> Detections:
|
||||
"""
|
||||
Performs Gaussian Soft Non-Maximum Suppression on detection set. Dispatch
|
||||
order: (1) if mask data present, IoU mask is used; (2) otherwise,
|
||||
axis-aligned box IoU is used. Oriented-box detections are not given
|
||||
dedicated OBB-IoU treatment and fall back to their axis-aligned `xyxy`.
|
||||
|
||||
Unlike `with_nms`, which discards overlapping detections outright,
|
||||
Soft-NMS keeps every detection and instead rescales its confidence by
|
||||
`score *= exp(-iou**2 / sigma)` for each higher-scoring, same-category
|
||||
overlap. By default (`score_threshold=None`) nothing is dropped — the
|
||||
method only rescales confidence, despite the "suppression" name; pass
|
||||
`score_threshold` to additionally filter the decayed scores into a real
|
||||
subset, matching `with_nms`'s return contract.
|
||||
|
||||
Args:
|
||||
sigma: Controls the strength of the confidence decay; must be
|
||||
greater than `0`. The lower the value the stronger the decay.
|
||||
No value of `sigma` reproduces `with_nms`'s hard cutoff —
|
||||
Soft-NMS never drops detections on its own.
|
||||
class_agnostic: Whether to perform class-agnostic Soft-NMS. If
|
||||
True, the class_id of each detection will be ignored.
|
||||
Defaults to False.
|
||||
score_threshold: If given, detections whose decayed confidence is
|
||||
at or below this value are dropped, producing a real subset
|
||||
(like `with_nms`). If `None` (default), all detections are
|
||||
kept, with their confidence rescaled in place on the returned
|
||||
copy.
|
||||
|
||||
Returns:
|
||||
A new Detections object with decayed confidence scores and,
|
||||
if `score_threshold` is given, filtered to a real subset.
|
||||
The original `Detections` instance is never modified.
|
||||
|
||||
Raises:
|
||||
ValueError: If `confidence` is None.
|
||||
If `class_id` is None and class_agnostic is False.
|
||||
If `sigma` is not greater than `0`.
|
||||
"""
|
||||
if len(self) == 0:
|
||||
return self
|
||||
|
||||
if self.confidence is None:
|
||||
raise ValueError(
|
||||
"Detections confidence must be given for Soft-NMS to be executed."
|
||||
)
|
||||
|
||||
predictions = self._build_nms_predictions(class_agnostic, "Soft-NMS")
|
||||
|
||||
if self.mask is not None:
|
||||
decayed_confidence = mask_soft_non_max_suppression(
|
||||
predictions=predictions,
|
||||
masks=self.mask,
|
||||
sigma=sigma,
|
||||
)
|
||||
else:
|
||||
decayed_confidence = box_soft_non_max_suppression(
|
||||
predictions=predictions,
|
||||
sigma=sigma,
|
||||
)
|
||||
|
||||
result = self.select(np.arange(len(self)))
|
||||
result.confidence = decayed_confidence
|
||||
|
||||
if score_threshold is None:
|
||||
return result
|
||||
return result.select(decayed_confidence > score_threshold)
|
||||
|
||||
def with_nmm(
|
||||
self,
|
||||
threshold: float = 0.5,
|
||||
|
|
@ -3107,28 +3192,7 @@ class Detections:
|
|||
"Detections confidence must be given for NMM to be executed."
|
||||
)
|
||||
|
||||
if class_agnostic:
|
||||
predictions = cast(
|
||||
npt.NDArray[np.floating],
|
||||
np.hstack((self.xyxy, self.confidence.reshape(-1, 1))),
|
||||
)
|
||||
else:
|
||||
if self.class_id is None:
|
||||
raise ValueError(
|
||||
"Detections class_id must be given for NMM to be executed. If "
|
||||
"you intended to perform class agnostic NMM "
|
||||
"set class_agnostic=True."
|
||||
)
|
||||
predictions = cast(
|
||||
npt.NDArray[np.floating],
|
||||
np.hstack(
|
||||
(
|
||||
self.xyxy,
|
||||
self.confidence.reshape(-1, 1),
|
||||
self.class_id.reshape(-1, 1),
|
||||
)
|
||||
),
|
||||
)
|
||||
predictions = self._build_nms_predictions(class_agnostic, "NMM")
|
||||
|
||||
if self.mask is not None:
|
||||
merge_groups = mask_non_max_merge(
|
||||
|
|
|
|||
|
|
@ -956,6 +956,91 @@ def mask_non_max_suppression(
|
|||
return cast(npt.NDArray[np.bool_], keep[sort_index.argsort()])
|
||||
|
||||
|
||||
def mask_soft_non_max_suppression(
|
||||
predictions: npt.NDArray[np.floating],
|
||||
masks: npt.NDArray[Any] | CompactMask,
|
||||
sigma: float = 0.5,
|
||||
overlap_metric: OverlapMetric = OverlapMetric.IOU,
|
||||
mask_dimension: int = 640,
|
||||
) -> npt.NDArray[np.floating]:
|
||||
"""
|
||||
Perform Soft Non-Maximum Suppression (Soft-NMS) on segmentation predictions.
|
||||
|
||||
Unlike `mask_non_max_suppression`, which discards overlapping masks outright,
|
||||
Soft-NMS keeps every detection and instead rescales its confidence by
|
||||
`score *= exp(-iou**2 / sigma)` for each higher-scoring, same-category
|
||||
overlap — the caller decides whether and where to threshold the result.
|
||||
A smaller `sigma` produces a stronger decay.
|
||||
|
||||
The 3rd positional parameter here is `sigma`, not `iou_threshold` as in
|
||||
`mask_non_max_suppression` — Soft-NMS has no threshold to suppress at, only
|
||||
a decay strength, so the two signatures intentionally diverge at that
|
||||
position.
|
||||
|
||||
IoU is computed exactly on the full-resolution masks for both dense and
|
||||
:class:`~supervision.detection.compact_mask.CompactMask` inputs. The
|
||||
`mask_dimension` parameter is kept for signature parity with
|
||||
`mask_non_max_suppression` but is not used — dense masks are **not** resized
|
||||
before IoU computation.
|
||||
|
||||
Args:
|
||||
predictions: A 2D array of object detection predictions in
|
||||
the format of `(x_min, y_min, x_max, y_max, score)`
|
||||
or `(x_min, y_min, x_max, y_max, score, class)`. Shape: `(N, 5)` or
|
||||
`(N, 6)`, where N is the number of predictions.
|
||||
masks: A 3D array of binary masks corresponding to the predictions.
|
||||
Shape: `(N, H, W)`, where N is the number of predictions, and H, W are the
|
||||
dimensions of each mask.
|
||||
sigma: Controls the strength of the confidence decay; must be greater
|
||||
than `0`. No value of `sigma` reproduces hard
|
||||
`mask_non_max_suppression` output — Soft-NMS never drops masks, only
|
||||
rescales confidence.
|
||||
overlap_metric: Metric used to compute the degree of overlap
|
||||
between pairs of masks (e.g., IoU, IoS).
|
||||
mask_dimension: Deprecated, unused. Kept for signature parity with
|
||||
`mask_non_max_suppression`.
|
||||
|
||||
Returns:
|
||||
An array containing the updated (decayed) confidence scores, in the
|
||||
same order as the input `predictions`.
|
||||
|
||||
Raises:
|
||||
ValueError: If `sigma` is not greater than `0`.
|
||||
|
||||
Examples:
|
||||
```pycon
|
||||
>>> import numpy as np
|
||||
>>> import supervision as sv
|
||||
>>> predictions = np.array([
|
||||
... [0, 0, 4, 4, 0.9, 0],
|
||||
... [0, 0, 4, 4, 0.8, 0],
|
||||
... ])
|
||||
>>> masks = np.zeros((2, 4, 4), dtype=bool)
|
||||
>>> masks[:, :2, :2] = True
|
||||
>>> sv.mask_soft_non_max_suppression(predictions, masks, sigma=0.5)
|
||||
array([0.9 , 0.10826823])
|
||||
|
||||
```
|
||||
"""
|
||||
_validate_sigma(sigma)
|
||||
rows, columns = predictions.shape
|
||||
|
||||
if columns == 5:
|
||||
predictions = np.c_[predictions, np.zeros(rows)]
|
||||
|
||||
sort_index = predictions[:, 4].argsort()[::-1]
|
||||
predictions = predictions[sort_index]
|
||||
masks = masks[sort_index]
|
||||
|
||||
ious = mask_iou_batch(masks, masks, overlap_metric)
|
||||
categories = predictions[:, 5]
|
||||
|
||||
decayed = _soft_nms_decay_from_iou_matrix(
|
||||
ious, categories, predictions[:, 4], sigma
|
||||
)
|
||||
return decayed[sort_index.argsort()]
|
||||
|
||||
|
||||
def _prepare_predictions_for_nms(
|
||||
predictions: npt.NDArray[np.floating],
|
||||
) -> tuple[npt.NDArray[np.int_], npt.NDArray[np.floating], npt.NDArray[np.floating]]:
|
||||
|
|
@ -995,6 +1080,36 @@ def _nms_loop_from_iou_matrix(
|
|||
return keep
|
||||
|
||||
|
||||
def _validate_sigma(sigma: float) -> None:
|
||||
"""Raise `ValueError` when a Soft-NMS `sigma` is not strictly positive."""
|
||||
if not sigma > 0:
|
||||
raise ValueError(f"Value of `sigma` must be greater than 0, {sigma} given.")
|
||||
|
||||
|
||||
def _soft_nms_decay_from_iou_matrix(
|
||||
ious: npt.NDArray[np.floating],
|
||||
categories: npt.NDArray[np.floating],
|
||||
scores: npt.NDArray[np.floating],
|
||||
sigma: float,
|
||||
) -> npt.NDArray[np.floating]:
|
||||
"""Vectorized Gaussian Soft-NMS confidence decay given a precomputed IoU matrix.
|
||||
|
||||
Assumes `ious`, `categories`, and `scores` are all sorted by descending score
|
||||
(as produced by `_prepare_predictions_for_nms`), and that `ious` is square with
|
||||
row/column order matching `categories`. Each detection's score is decayed once
|
||||
per higher-scoring, same-category detection that precedes it — equivalent to
|
||||
the reference single-pass (no re-sort) Soft-NMS loop, computed as a single
|
||||
vectorized closed form: `decayed[j] = scores[j] * exp(-sum_i(ious[i, j]**2) /
|
||||
sigma)` summed over same-category `i < j`.
|
||||
"""
|
||||
rows = len(ious)
|
||||
same_category = categories[:, None] == categories[None, :]
|
||||
precedes = np.triu(np.ones((rows, rows), dtype=bool), k=1)
|
||||
weighted_iou_sq = np.where(precedes & same_category, ious**2, 0.0)
|
||||
decay_exponent = weighted_iou_sq.sum(axis=0)
|
||||
return cast(npt.NDArray[np.floating], scores * np.exp(-decay_exponent / sigma))
|
||||
|
||||
|
||||
def box_non_max_suppression(
|
||||
predictions: npt.NDArray[np.floating],
|
||||
iou_threshold: float = 0.5,
|
||||
|
|
@ -1041,6 +1156,61 @@ def box_non_max_suppression(
|
|||
return result
|
||||
|
||||
|
||||
def box_soft_non_max_suppression(
|
||||
predictions: npt.NDArray[np.floating],
|
||||
sigma: float = 0.5,
|
||||
overlap_metric: OverlapMetric = OverlapMetric.IOU,
|
||||
) -> npt.NDArray[np.floating]:
|
||||
"""
|
||||
Perform Soft Non-Maximum Suppression (Soft-NMS) on object detection predictions.
|
||||
|
||||
Unlike `box_non_max_suppression`, which discards overlapping boxes outright,
|
||||
Soft-NMS keeps every detection and instead rescales its confidence by
|
||||
`score *= exp(-iou**2 / sigma)` for each higher-scoring, same-category
|
||||
overlap — the caller decides whether and where to threshold the result.
|
||||
A smaller `sigma` produces a stronger decay.
|
||||
|
||||
Args:
|
||||
predictions: An array of object detection predictions in
|
||||
the format of `(x_min, y_min, x_max, y_max, score)`
|
||||
or `(x_min, y_min, x_max, y_max, score, class)`.
|
||||
sigma: Controls the strength of the confidence decay; must be greater
|
||||
than `0`. No value of `sigma` reproduces hard
|
||||
`box_non_max_suppression` output — Soft-NMS never drops boxes, only
|
||||
rescales confidence.
|
||||
overlap_metric: Metric used to compute the degree of overlap
|
||||
between pairs of boxes (e.g., IoU, IoS).
|
||||
|
||||
Returns:
|
||||
An array containing the updated (decayed) confidence scores, in the
|
||||
same order as the input `predictions`.
|
||||
|
||||
Raises:
|
||||
ValueError: If `sigma` is not greater than `0`.
|
||||
|
||||
Examples:
|
||||
```pycon
|
||||
>>> import numpy as np
|
||||
>>> import supervision as sv
|
||||
>>> predictions = np.array([
|
||||
... [0, 0, 4, 4, 0.9, 0],
|
||||
... [0, 0, 4, 4, 0.8, 0],
|
||||
... ])
|
||||
>>> sv.box_soft_non_max_suppression(predictions, sigma=0.5)
|
||||
array([0.9 , 0.10826823])
|
||||
|
||||
```
|
||||
"""
|
||||
_validate_sigma(sigma)
|
||||
sort_index, predictions, categories = _prepare_predictions_for_nms(predictions)
|
||||
ious = box_iou_batch(predictions[:, :4], predictions[:, :4], overlap_metric)
|
||||
decayed = _soft_nms_decay_from_iou_matrix(
|
||||
ious, categories, predictions[:, 4], sigma
|
||||
)
|
||||
result_scores: npt.NDArray[np.floating] = decayed[sort_index.argsort()]
|
||||
return result_scores
|
||||
|
||||
|
||||
def _group_overlapping_masks(
|
||||
predictions: npt.NDArray[np.floating],
|
||||
masks: npt.NDArray[np.bool_],
|
||||
|
|
|
|||
|
|
@ -1669,6 +1669,130 @@ class TestDetectionsOverlapValidation:
|
|||
getattr(detections, method)(threshold=0.5)
|
||||
|
||||
|
||||
class TestDetectionsWithSoftNms:
|
||||
"""`Detections.with_soft_nms` — decay semantics, dispatch, non-mutation."""
|
||||
|
||||
def test_requires_confidence(self) -> None:
|
||||
"""Missing confidence raises a descriptive `ValueError`."""
|
||||
detections = Detections(
|
||||
xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Detections confidence must be given"):
|
||||
detections.with_soft_nms()
|
||||
|
||||
def test_requires_class_id_when_not_class_agnostic(self) -> None:
|
||||
"""Missing class IDs raise a descriptive `ValueError`."""
|
||||
detections = Detections(
|
||||
xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
|
||||
confidence=np.array([0.9], dtype=np.float32),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Detections class_id must be given"):
|
||||
detections.with_soft_nms()
|
||||
|
||||
def test_rejects_non_positive_sigma(self) -> None:
|
||||
"""`sigma` must be strictly greater than 0."""
|
||||
detections = Detections(
|
||||
xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
|
||||
confidence=np.array([0.9], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="sigma"):
|
||||
detections.with_soft_nms(sigma=0)
|
||||
|
||||
def test_returns_self_for_empty_detections(self) -> None:
|
||||
"""Empty input short-circuits without dispatching to Soft-NMS."""
|
||||
detections = Detections.empty()
|
||||
|
||||
result = detections.with_soft_nms()
|
||||
|
||||
assert result is detections
|
||||
|
||||
def test_does_not_mutate_original(self) -> None:
|
||||
"""The source `Detections` confidence is untouched after the call."""
|
||||
detections = Detections(
|
||||
xyxy=np.array([[10, 10, 40, 40], [15, 15, 40, 40]], dtype=np.float32),
|
||||
confidence=np.array([0.8, 0.9], dtype=np.float32),
|
||||
class_id=np.array([0, 0], dtype=int),
|
||||
)
|
||||
original_confidence = detections.confidence.copy()
|
||||
|
||||
result = detections.with_soft_nms(sigma=0.2)
|
||||
|
||||
assert result is not detections
|
||||
np.testing.assert_array_equal(detections.confidence, original_confidence)
|
||||
np.testing.assert_almost_equal(result.confidence, [0.07176137, 0.9], decimal=5)
|
||||
|
||||
def test_box_path_class_agnostic(self) -> None:
|
||||
"""`class_agnostic=True` skips the class_id requirement for boxes."""
|
||||
detections = Detections(
|
||||
xyxy=np.array([[10, 10, 40, 40], [15, 15, 40, 40]], dtype=np.float32),
|
||||
confidence=np.array([0.8, 0.9], dtype=np.float32),
|
||||
)
|
||||
|
||||
result = detections.with_soft_nms(sigma=0.2, class_agnostic=True)
|
||||
|
||||
np.testing.assert_almost_equal(result.confidence, [0.07176137, 0.9], decimal=5)
|
||||
|
||||
def test_mask_path_dispatch(self) -> None:
|
||||
"""Mask data present routes through `mask_soft_non_max_suppression`."""
|
||||
masks = np.zeros((2, 4, 4), dtype=bool)
|
||||
masks[:, :2, :2] = True
|
||||
detections = Detections(
|
||||
xyxy=np.array([[0, 0, 4, 4], [0, 0, 4, 4]], dtype=np.float32),
|
||||
confidence=np.array([0.9, 0.8], dtype=np.float32),
|
||||
class_id=np.array([0, 0], dtype=int),
|
||||
mask=masks,
|
||||
)
|
||||
|
||||
result = detections.with_soft_nms(sigma=0.5)
|
||||
|
||||
np.testing.assert_almost_equal(result.confidence, [0.9, 0.10826823], decimal=5)
|
||||
assert result.mask is not detections.mask
|
||||
|
||||
def test_single_mask_self_decay_class_agnostic(self) -> None:
|
||||
"""A single class-agnostic mask detection is left undecayed (no peer)."""
|
||||
masks = np.zeros((1, 4, 4), dtype=bool)
|
||||
masks[:, :2, :2] = True
|
||||
detections = Detections(
|
||||
xyxy=np.array([[0, 0, 4, 4]], dtype=np.float32),
|
||||
confidence=np.array([0.8], dtype=np.float32),
|
||||
mask=masks,
|
||||
)
|
||||
|
||||
result = detections.with_soft_nms(sigma=0.5, class_agnostic=True)
|
||||
|
||||
np.testing.assert_almost_equal(result.confidence, [0.8], decimal=5)
|
||||
|
||||
def test_score_threshold_filters_to_subset(self) -> None:
|
||||
"""`score_threshold` drops decayed detections, producing a real subset."""
|
||||
detections = Detections(
|
||||
xyxy=np.array([[10, 10, 40, 40], [15, 15, 40, 40]], dtype=np.float32),
|
||||
confidence=np.array([0.8, 0.9], dtype=np.float32),
|
||||
class_id=np.array([0, 0], dtype=int),
|
||||
)
|
||||
|
||||
result = detections.with_soft_nms(sigma=0.2, score_threshold=0.5)
|
||||
|
||||
assert len(result) == 1
|
||||
np.testing.assert_almost_equal(result.confidence, [0.9], decimal=5)
|
||||
|
||||
def test_none_score_threshold_keeps_all_detections(self) -> None:
|
||||
"""Default `score_threshold=None` never drops a detection."""
|
||||
detections = Detections(
|
||||
xyxy=np.array([[10, 10, 40, 40], [15, 15, 40, 40]], dtype=np.float32),
|
||||
confidence=np.array([0.8, 0.9], dtype=np.float32),
|
||||
class_id=np.array([0, 0], dtype=int),
|
||||
)
|
||||
|
||||
result = detections.with_soft_nms(sigma=0.2)
|
||||
|
||||
assert len(result) == 2
|
||||
|
||||
|
||||
class TestGetAnchorsObbDispatch:
|
||||
"""`get_anchors_coordinates` reads oriented corners when OBB data is present."""
|
||||
|
||||
|
|
|
|||
|
|
@ -14,9 +14,11 @@ from supervision.detection.utils.iou_and_nms import (
|
|||
box_iou_batch_with_jaccard,
|
||||
box_non_max_merge,
|
||||
box_non_max_suppression,
|
||||
box_soft_non_max_suppression,
|
||||
mask_iou_batch,
|
||||
mask_non_max_merge,
|
||||
mask_non_max_suppression,
|
||||
mask_soft_non_max_suppression,
|
||||
oriented_box_iou_batch,
|
||||
oriented_box_non_max_merge,
|
||||
oriented_box_non_max_suppression,
|
||||
|
|
@ -283,6 +285,139 @@ def test_box_non_max_suppression(
|
|||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("predictions", "sigma", "expected_result", "exception"),
|
||||
[
|
||||
(
|
||||
np.empty(shape=(0, 5)),
|
||||
0.1,
|
||||
np.array([]),
|
||||
DoesNotRaise(),
|
||||
), # empty predictions
|
||||
(
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.8]]),
|
||||
0.8,
|
||||
np.array([0.8]),
|
||||
DoesNotRaise(),
|
||||
), # single box with no category
|
||||
(
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.8, 0]]),
|
||||
0.9,
|
||||
np.array([0.8]),
|
||||
DoesNotRaise(),
|
||||
), # single box with category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9],
|
||||
]
|
||||
),
|
||||
0.2,
|
||||
np.array([0.07176137, 0.9]),
|
||||
DoesNotRaise(),
|
||||
), # two boxes with no category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9, 1],
|
||||
]
|
||||
),
|
||||
0.3,
|
||||
np.array([0.8, 0.9]),
|
||||
DoesNotRaise(),
|
||||
), # two boxes with different category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9, 0],
|
||||
]
|
||||
),
|
||||
0.9,
|
||||
np.array([0.46814354, 0.9]),
|
||||
DoesNotRaise(),
|
||||
), # two boxes with same category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85],
|
||||
]
|
||||
),
|
||||
0.7,
|
||||
np.array([0.42648529, 0.9, 0.53109062]),
|
||||
DoesNotRaise(),
|
||||
), # three boxes with no category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9, 1],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85, 2],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([0.8, 0.9, 0.85]),
|
||||
DoesNotRaise(),
|
||||
), # three boxes with different category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9, 0],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85, 1],
|
||||
]
|
||||
),
|
||||
0.9,
|
||||
np.array([0.55491779, 0.9, 0.85]),
|
||||
DoesNotRaise(),
|
||||
), # three boxes with different category, one class isolated
|
||||
(
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.5]]),
|
||||
0.0,
|
||||
None,
|
||||
pytest.raises(ValueError, match="sigma"),
|
||||
), # sigma at the lower boundary (not > 0) must raise
|
||||
(
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.5]]),
|
||||
-0.1,
|
||||
None,
|
||||
pytest.raises(ValueError, match="sigma"),
|
||||
), # negative sigma must raise
|
||||
(
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.5]]),
|
||||
5.0,
|
||||
np.array([0.5]),
|
||||
DoesNotRaise(),
|
||||
), # sigma > 1 is valid for Soft-NMS (no upper bound)
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[10.0, 10.0, 40.0, 40.0, 0.7],
|
||||
[10.0, 10.0, 40.0, 40.0, 0.7],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([0.0947347, 0.7]),
|
||||
DoesNotRaise(),
|
||||
), # tied confidence: argsort's tie order picks a winner, exactly one of
|
||||
# the two identical, fully-overlapping boxes is decayed
|
||||
],
|
||||
)
|
||||
def test_box_soft_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
sigma: float,
|
||||
expected_result: np.ndarray | None,
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = box_soft_non_max_suppression(predictions=predictions, sigma=sigma)
|
||||
np.testing.assert_almost_equal(result, expected_result, decimal=5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("predictions", "masks", "iou_threshold", "expected_result", "exception"),
|
||||
[
|
||||
|
|
@ -489,6 +624,157 @@ def test_mask_non_max_suppression(
|
|||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("predictions", "masks", "sigma", "expected_result", "exception"),
|
||||
[
|
||||
(
|
||||
np.empty((0, 6)),
|
||||
np.empty((0, 5, 5)),
|
||||
0.1,
|
||||
np.array([]),
|
||||
DoesNotRaise(),
|
||||
), # empty predictions and masks
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
]
|
||||
]
|
||||
),
|
||||
0.2,
|
||||
np.array([0.8]),
|
||||
DoesNotRaise(),
|
||||
), # single mask with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8, 0]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
]
|
||||
]
|
||||
),
|
||||
0.99,
|
||||
np.array([0.8]),
|
||||
DoesNotRaise(),
|
||||
), # single mask with category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8], [0, 0, 0, 0, 0.9]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.8,
|
||||
np.array([0.8, 0.9]),
|
||||
DoesNotRaise(),
|
||||
), # two masks non-overlapping with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8], [0, 0, 0, 0, 0.9]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.6,
|
||||
np.array([0.5273925, 0.9]),
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with no category — exact
|
||||
# full-resolution IoU (this module no longer resizes masks before
|
||||
# computing IoU, matching mask_non_max_suppression's convention)
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8, 0], [0, 0, 0, 0, 0.9, 1]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.9,
|
||||
np.array([0.8, 0.9]),
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with different category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.5]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
]
|
||||
]
|
||||
),
|
||||
0.0,
|
||||
None,
|
||||
pytest.raises(ValueError, match="sigma"),
|
||||
), # sigma at the lower boundary (not > 0) must raise
|
||||
],
|
||||
)
|
||||
def test_mask_soft_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
sigma: float,
|
||||
expected_result: np.ndarray | None,
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = mask_soft_non_max_suppression(
|
||||
predictions=predictions, masks=masks, sigma=sigma
|
||||
)
|
||||
np.testing.assert_almost_equal(result, expected_result, decimal=5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("predictions", "masks", "iou_threshold", "expected_result", "exception"),
|
||||
[
|
||||
|
|
|
|||
Loading…
Reference in New Issue