perf: give NMS and greedy NMM back the stored-pair path they lost in 0.12.5 (#1443)

Author: onuralpszr

sahi/postprocess/_numpy_backend.py

@@ -12,11 +12,14 @@
 from sahi.postprocess._sparse_backend import (
     _safe_ratio,
     build_sparse_matches,
+    greedy_nmm_sparse,
     greedy_nmm_streaming,
     nmm_sparse,
     nmm_streaming,
+    nms_sparse,
     nms_streaming,
     should_stream_nmm,
+    should_stream_survivor,
     should_use_sparse,
 )
 
@@ -400,8 +403,11 @@ def nms_numpy(
     if matches_all_pairs(match_threshold):
         return nms_match_all(predictions)
     if should_use_sparse(len(predictions), match_threshold):
-        boxes, areas, sorted_idxs = _prepare_streaming(predictions)
-        return nms_streaming(boxes, areas, match_metric, match_threshold, sorted_idxs)
+        if should_stream_survivor(predictions[:, :4]):
+            boxes, areas, sorted_idxs = _prepare_streaming(predictions)
+            return nms_streaming(boxes, areas, match_metric, match_threshold, sorted_idxs)
+        indptr, indices, sorted_idxs = _prepare_sparse(predictions, match_metric, match_threshold)
+        return nms_sparse(indptr, indices, sorted_idxs)
     matrix, sorted_idxs = _prepare_matrix(predictions, match_metric)
     return nms_from_matrix(matrix, sorted_idxs, match_threshold)
 
@@ -415,8 +421,11 @@ def greedy_nmm_numpy(
     if matches_all_pairs(match_threshold):
         return greedy_nmm_match_all(predictions)
     if should_use_sparse(len(predictions), match_threshold):
-        boxes, areas, sorted_idxs = _prepare_streaming(predictions)
-        return greedy_nmm_streaming(boxes, areas, match_metric, match_threshold, sorted_idxs)
+        if should_stream_survivor(predictions[:, :4]):
+            boxes, areas, sorted_idxs = _prepare_streaming(predictions)
+            return greedy_nmm_streaming(boxes, areas, match_metric, match_threshold, sorted_idxs)
+        indptr, indices, sorted_idxs = _prepare_sparse(predictions, match_metric, match_threshold)
+        return greedy_nmm_sparse(indptr, indices, sorted_idxs)
     matrix, sorted_idxs = _prepare_matrix(predictions, match_metric)
     return greedy_nmm_from_matrix(matrix, sorted_idxs, match_threshold)
 

sahi/postprocess/_sparse_backend.py

@@ -32,6 +32,12 @@
 # crossover barely moves with N, so it is a degree and not a pair budget.
 NMM_MAX_DEGREE = 40.0
 
+# NMS and greedy NMM only ever read rows of boxes that survived, so they settle
+# far fewer rows than NMM and the stored list keeps paying off much longer. Both
+# measure the same crossover, around degree 175; staying under it keeps the pair
+# list bounded on the crowded inputs streaming exists to protect.
+SURVIVOR_MAX_DEGREE = 150.0
+
 # Probing the tree in one call would return sample x degree pairs at once, which
 # on crowded boxes dwarfs everything the path being chosen goes on to allocate.
 # Only the count is wanted, so the probe is spent a chunk at a time.
@@ -96,17 +102,31 @@ def _dominates(left: int | np.ndarray, right: np.ndarray, scores: np.ndarray, bo
     return lower_score | (score_equal & ~left_lt)
 
 
-def should_stream_nmm(boxes: np.ndarray) -> bool:
-    """Return whether NMM should answer rows from the tree instead of storing pairs.
+def should_stream(boxes: np.ndarray, max_degree: float) -> bool:
+    """Return whether to answer rows from the tree instead of storing pairs.
 
     Args:
         boxes: Array of shape (N, 4) with columns [x1, y1, x2, y2].
+        max_degree: Mean degree above which streaming wins, per caller.
 
     Returns:
         True when the boxes intersect each other often enough that reading rows
         from the tree beats storing them.
     """
-    return estimate_mean_degree(boxes) > NMM_MAX_DEGREE
+    return estimate_mean_degree(boxes) > max_degree
+
+
+def should_stream_nmm(boxes: np.ndarray) -> bool:
+    """Return whether NMM should answer rows from the tree instead of storing pairs."""
+    return should_stream(boxes, NMM_MAX_DEGREE)
+
+
+def should_stream_survivor(boxes: np.ndarray) -> bool:
+    """Return whether NMS and greedy NMM should read rows from the tree.
+
+    Both settle a row only for a box that survived, so they share a crossover.
+    """
+    return should_stream(boxes, SURVIVOR_MAX_DEGREE)
 
 
 def should_use_sparse(n: int, match_threshold: float) -> bool:
@@ -383,6 +403,70 @@ def _dominates_all(indptr: np.ndarray, indices: np.ndarray, scores: np.ndarray,
     return _dominates(rows, indices, scores, boxes)
 
 
+def nms_sparse(indptr: np.ndarray, indices: np.ndarray, sorted_idxs: np.ndarray) -> list[int]:
+    """NMS over a CSR match adjacency. Mirrors ``nms_from_matrix``.
+
+    Suppressing a whole row is one vectorized store, so while the pair list
+    stays bounded this beats querying the tree per survivor.
+
+    Args:
+        indptr: CSR row pointers of length N + 1.
+        indices: CSR column indices.
+        sorted_idxs: Indices sorted by score descending.
+
+    Returns:
+        List of kept indices sorted by score descending.
+    """
+    keep: list[int] = []
+    suppressed = np.zeros(len(indptr) - 1, dtype=bool)
+
+    for idx in sorted_idxs:
+        if suppressed[idx]:
+            continue
+        keep.append(int(idx))
+        suppressed[indices[indptr[idx] : indptr[idx + 1]]] = True
+
+    return keep
+
+
+def greedy_nmm_sparse(
+    indptr: np.ndarray,
+    indices: np.ndarray,
+    sorted_idxs: np.ndarray,
+) -> dict[int, list[int]]:
+    """Greedy NMM over a CSR match adjacency. Mirrors ``greedy_nmm_from_matrix``.
+
+    Args:
+        indptr: CSR row pointers of length N + 1.
+        indices: CSR column indices.
+        sorted_idxs: Indices sorted by score descending.
+
+    Returns:
+        Dict mapping each kept index to a list of indices merged into it.
+    """
+    n = len(indptr) - 1
+    suppressed = np.zeros(n, dtype=bool)
+
+    # The dense loop only considers candidates that come later in score order,
+    # and emits them in that order.
+    rank = np.empty(n, dtype=np.intp)
+    rank[sorted_idxs] = np.arange(n)
+
+    keep_to_merge_list: dict[int, list[int]] = {}
+    for position, idx in enumerate(sorted_idxs):
+        if suppressed[idx]:
+            continue
+
+        neighbours = indices[indptr[idx] : indptr[idx + 1]]
+        merge_indices = neighbours[(rank[neighbours] > position) & ~suppressed[neighbours]]
+        merge_indices = merge_indices[np.argsort(rank[merge_indices])]
+
+        suppressed[merge_indices] = True
+        keep_to_merge_list[int(idx)] = merge_indices.tolist()
+
+    return keep_to_merge_list
+
+
 def nmm_sparse(
     indptr: np.ndarray,
     indices: np.ndarray,

tests/test_sparse_backend.py

@@ -19,11 +19,14 @@
     SPARSE_MIN_BOXES,
     MatchQuery,
     build_sparse_matches,
+    greedy_nmm_sparse,
     greedy_nmm_streaming,
     nmm_sparse,
     nmm_streaming,
+    nms_sparse,
     nms_streaming,
     should_stream_nmm,
+    should_stream_survivor,
     should_use_sparse,
 )
 
@@ -109,10 +112,12 @@ def test_dense_csr_and_streaming_agree(match_metric: str, match_threshold: float
     case = _case(predictions, match_metric, match_threshold)
     streaming_args = (case.boxes, case.areas, match_metric, match_threshold, case.sorted_idxs)
 
-    assert nms_streaming(*streaming_args) == nms_from_matrix(case.matrix, case.sorted_idxs, match_threshold)
-    assert greedy_nmm_streaming(*streaming_args) == greedy_nmm_from_matrix(
-        case.matrix, case.sorted_idxs, match_threshold
-    )
+    dense_nms = nms_from_matrix(case.matrix, case.sorted_idxs, match_threshold)
+    assert nms_streaming(*streaming_args) == dense_nms
+    assert nms_sparse(case.indptr, case.indices, case.sorted_idxs) == dense_nms
+    dense_greedy = greedy_nmm_from_matrix(case.matrix, case.sorted_idxs, match_threshold)
+    assert greedy_nmm_streaming(*streaming_args) == dense_greedy
+    assert greedy_nmm_sparse(case.indptr, case.indices, case.sorted_idxs) == dense_greedy
 
     dense_nmm = nmm_from_matrix(case.matrix, case.sorted_idxs, case.scores, case.boxes, match_threshold)
     assert nmm_sparse(case.indptr, case.indices, case.sorted_idxs, case.scores, case.boxes) == dense_nmm
@@ -138,7 +143,13 @@ def test_streaming_agrees_with_dense_once_the_tree_is_rebuilt(match_metric: str)
     streaming_args = (case.boxes, case.areas, match_metric, 0.3, case.sorted_idxs)
 
     assert nms_streaming(*streaming_args) == nms_from_matrix(case.matrix, case.sorted_idxs, 0.3)
+    assert nms_sparse(case.indptr, case.indices, case.sorted_idxs) == nms_from_matrix(
+        case.matrix, case.sorted_idxs, 0.3
+    )
     assert greedy_nmm_streaming(*streaming_args) == greedy_nmm_from_matrix(case.matrix, case.sorted_idxs, 0.3)
+    assert greedy_nmm_sparse(case.indptr, case.indices, case.sorted_idxs) == greedy_nmm_from_matrix(
+        case.matrix, case.sorted_idxs, 0.3
+    )
 
     result = nmm_streaming(*streaming_args, case.scores)
     assert result == nmm_from_matrix(case.matrix, case.sorted_idxs, case.scores, case.boxes, 0.3)
@@ -177,6 +188,11 @@ def test_nmm_streams_only_when_boxes_are_crowded() -> None:
     assert should_stream_nmm(scattered[:, :4]) is False
     assert should_stream_nmm(crowded[:, :4]) is True
 
+    # NMS and greedy NMM settle fewer rows than NMM, so they keep the stored
+    # pairs for longer, but the same crowding must still push them to streaming.
+    assert should_stream_survivor(scattered[:, :4]) is False
+    assert should_stream_survivor(crowded[:, :4]) is True
+
     # Both routes must agree with the stored-pair result and still group every box.
     for name, predictions in (("scattered", scattered), ("crowded", crowded)):
         boxes, scores = predictions[:, :4], predictions[:, 4]