From 89d49c2e9342b2eabe73fced50074c2993bf9c32 Mon Sep 17 00:00:00 2001
From: HALLOUARD <57447861+YHallouard@users.noreply.github.com>
Date: Tue, 21 Jul 2026 10:25:14 +0200
Subject: [PATCH] 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>
---
docs/detection/utils/iou_and_nms.md | 12 +
src/supervision/__init__.py | 4 +
src/supervision/detection/core.py | 152 +++++++---
.../detection/utils/iou_and_nms.py | 170 +++++++++++
tests/detection/test_core.py | 124 ++++++++
tests/detection/utils/test_iou_and_nms.py | 286 ++++++++++++++++++
6 files changed, 704 insertions(+), 44 deletions(-)
diff --git a/docs/detection/utils/iou_and_nms.md b/docs/detection/utils/iou_and_nms.md
index 93e98e8c..49dc1699 100644
--- a/docs/detection/utils/iou_and_nms.md
+++ b/docs/detection/utils/iou_and_nms.md
@@ -52,12 +52,24 @@ comments: true
:::supervision.detection.utils.iou_and_nms.box_non_max_suppression
+
+
+:::supervision.detection.utils.iou_and_nms.box_soft_non_max_suppression
+
:::supervision.detection.utils.iou_and_nms.mask_non_max_suppression
+
+
+:::supervision.detection.utils.iou_and_nms.mask_soft_non_max_suppression
+
diff --git a/src/supervision/__init__.py b/src/supervision/__init__.py
index fbe75c0d..f01cc058 100644
--- a/src/supervision/__init__.py
+++ b/src/supervision/__init__.py
@@ -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",
diff --git a/src/supervision/detection/core.py b/src/supervision/detection/core.py
index d9f4c835..2e2df10a 100644
--- a/src/supervision/detection/core.py
+++ b/src/supervision/detection/core.py
@@ -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(
diff --git a/src/supervision/detection/utils/iou_and_nms.py b/src/supervision/detection/utils/iou_and_nms.py
index 6ea01fe3..3b9876b7 100644
--- a/src/supervision/detection/utils/iou_and_nms.py
+++ b/src/supervision/detection/utils/iou_and_nms.py
@@ -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_],
diff --git a/tests/detection/test_core.py b/tests/detection/test_core.py
index d5083207..858a1cee 100644
--- a/tests/detection/test_core.py
+++ b/tests/detection/test_core.py
@@ -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."""
diff --git a/tests/detection/utils/test_iou_and_nms.py b/tests/detection/utils/test_iou_and_nms.py
index c4c0e870..e4b227cf 100644
--- a/tests/detection/utils/test_iou_and_nms.py
+++ b/tests/detection/utils/test_iou_and_nms.py
@@ -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"),
[