Better edge case handling

Author: remydubois

changelog.md

@@ -1,6 +1,11 @@
 Changelog
 =========
 
+Version 0.3.1
+------------
+- Edge case where all box scores are zero (or all below threshold) is now handled (threw uggly error before)
+
+
 Version 0.3.0
 ------------
 - Added multiclass support for NMS

lsnms/nms.py

@@ -22,6 +22,9 @@ def _nms(
     # Discard boxes below score threshold right now to avoid building the tree on useless boxes
     boxes = boxes[scores > score_threshold]
 
+    if len(boxes) == 0:
+        return np.zeros(0, dtype=np.int64)
+
     # Build the BallTree
     rtree = RNode(boxes, tree_leaf_size, max_spread_axis(boxes), None)
     rtree.build()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "lsnms"
-version = "0.3.0"
+version = "0.3.1"
 description = "Large Scale Non Maximum Suppression"
 authors = ["Rémy Dubois <remydubois14@gmail.com>"]
 license = "MIT"

tests/test_nms.py

@@ -29,6 +29,14 @@ def test_rtree_nms():
     assert np.allclose(k1, k2)
 
 
+def test_empty_nms():
+    boxes, scores = datagen()
+    # Put all the scores to zero artificially
+    keep = nms(boxes, scores * 0.)
+
+    assert keep.size == 0
+    
+
 def test_rtree_multiclass_nms():
 
     boxes, scores = datagen()