feat: Linear sum assignment

bindings/python/powerboxes/__init__.py

@@ -18,6 +18,7 @@
     _dtype_to_func_rtree_nms,
     _dtype_to_func_rtree_rotated_nms,
     _dtype_to_func_tiou_distance,
+    _dtype_to_func_lsap_iou,
 )
 from ._powerboxes import draw_boxes as _draw_boxes
 from ._powerboxes import draw_rotated_boxes as _draw_rotated_boxes
@@ -646,6 +647,25 @@ def draw_rotated_boxes(
     return _draw_rotated_boxes(image, boxes, colors, thickness, filled, opacity)
 
 
+def lsap_iou(
+    boxes1: npt.NDArray[T], boxes2: npt.NDArray[T], iou_threshold: float = 0.0
+) -> list[tuple[int]]:
+    """Perform optimal asssignement between 2 sets of axis-aligned boxes, based on their IoU.
+
+    Uses Shortest Augmenting Path Algorithm.
+
+    Args:
+        boxes1: first set of boxes with shape (N, 4)
+        boxes2: second set of boxes with shape (N, 4)
+        iou_threshold: threshold below which two boxes cannot get matched together.
+
+    Returns:
+        list of assignments. An assignment is a tuple (x, y) where x is the indice in boxes1 and
+        y is the matching indice in boxes2.
+    """
+    return _dispatch2(_dtype_to_func_lsap_iou, boxes1, boxes2, iou_threshold)
+
+
 __all__ = [
     "ciou_distance",
     "diou_distance",
@@ -670,6 +690,7 @@ def draw_rotated_boxes(
     "rtree_rotated_nms",
     "draw_boxes",
     "draw_rotated_boxes",
+    "lsap_iou",
     "supported_dtypes",
     "__version__",
 ]

bindings/python/powerboxes/_dispatch.py

@@ -40,3 +40,4 @@ def _build_dispatch(prefix: str, suffixes=_ALL_SUFFIXES) -> dict:
 _dtype_to_func_rtree_rotated_nms = _build_dispatch(
     "rtree_rotated_nms", _SIGNED_SUFFIXES
 )
+_dtype_to_func_lsap_iou = _build_dispatch("lsap_iou")

bindings/src/lib.rs

@@ -5,7 +5,7 @@ use std::fmt::Debug;
 use ndarray::Array1;
 use num_traits::{Bounded, Float, Num, Signed, ToPrimitive};
 use numpy::{PyArray1, PyArray2, PyArray3, PyArrayMethods};
-use powerboxesrs::{boxes, ciou, diou, draw, giou, iou, nms, tiou};
+use powerboxesrs::{assignments, boxes, ciou, diou, draw, giou, iou, nms, tiou};
 use pyo3::prelude::*;
 use utils::{preprocess_array1, preprocess_array3, preprocess_boxes, preprocess_rotated_boxes};
 
@@ -68,6 +68,22 @@ macro_rules! impl_distance2_fn {
     };
 }
 
+/// Generate a typed `#[pyfunction]` for `(py, boxes1, boxes2, iou_threshold) -> Vec<(usize, usize)>`.
+macro_rules! impl_assignment_fn {
+    ($prefix:ident, $generic:ident, $T:ty, $suffix:ident) => {
+        ::paste::paste! {
+            #[pyfunction]
+            fn [<$prefix _ $suffix>](
+                boxes1: &Bound<'_, PyArray2<$T>>,
+                boxes2: &Bound<'_, PyArray2<$T>>,
+                iou_threshold: f64,
+            ) -> PyResult<Vec<(usize, usize)>> {
+                $generic(boxes1, boxes2, iou_threshold)
+            }
+        }
+    };
+}
+
 /// Generate a typed `#[pyfunction]` for `(py, boxes) -> Array1<f64>`.
 macro_rules! impl_unary_f64_fn {
     ($prefix:ident, $generic:ident, $T:ty, $suffix:ident) => {
@@ -208,6 +224,8 @@ fn _powerboxes(m: &Bound<'_, PyModule>) -> PyResult<()> {
     register_typed!(m, rtree_nms, [f64, f32, i64, i32, i16]);
     // Rtree Rotated NMS (signed + float only)
     register_typed!(m, rtree_rotated_nms, [f64, f32, i64, i32, i16]);
+    // LSAP on IoU
+    register_typed!(m, lsap_iou, [f64, f32, i64, i32, i16, u64, u32, u16, u8]);
     // Masks to boxes
     m.add_function(wrap_pyfunction!(masks_to_boxes, m)?)?;
     // Rotated IoU
@@ -798,3 +816,19 @@ for_each_numeric_type!(
     rtree_rotated_nms_generic,
     signed
 );
+
+// Linear Sum Assignments IoU
+fn lsap_iou_generic<T>(
+    boxes1: &Bound<'_, PyArray2<T>>,
+    boxes2: &Bound<'_, PyArray2<T>>,
+    iou_threshold: f64,
+) -> PyResult<Vec<(usize, usize)>>
+where
+    T: Num + ToPrimitive + PartialOrd + numpy::Element + Copy,
+{
+    let boxes1 = preprocess_boxes(boxes1)?;
+    let boxes2 = preprocess_boxes(boxes2)?;
+    let asgmt = assignments::lsap_iou(boxes1, boxes2, iou_threshold);
+    Ok(asgmt)
+}
+for_each_numeric_type!(impl_assignment_fn, lsap_iou, lsap_iou_generic);

bindings/tests/test_dtypes.py

@@ -9,6 +9,7 @@
     diou_distance,
     giou_distance,
     iou_distance,
+    lsap_iou,
     masks_to_boxes,
     nms,
     parallel_giou_distance,
@@ -472,3 +473,50 @@ def test_rotated_nms_bad_dtype():
     scores = np.random.random((100,))
     with pytest.raises(TypeError):
         rotated_nms(boxes.astype(unsuported_dtype_example), scores, 0.5, 0.5)
+
+
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_lsap_iou(dtype):
+    # seeded + tight bounds so the uint8 cast is valid AND box_areas's
+    # native-dtype multiply (x2-x1)*(y2-y1) doesn't overflow u8 (needs wh<=15).
+    rng = np.random.default_rng(0)
+    topleft = rng.uniform(0, 200, size=(50, 2))
+    wh = rng.uniform(5, 15, size=(50, 2))
+    boxes1 = np.concatenate([topleft, topleft + wh], axis=1)
+    topleft2 = rng.uniform(0, 200, size=(50, 2))
+    boxes2 = np.concatenate([topleft2, topleft2 + wh], axis=1)
+    pairs = lsap_iou(boxes1.astype(dtype), boxes2.astype(dtype))
+    assert isinstance(pairs, list)
+    for p in pairs:
+        assert len(p) == 2
+        assert 0 <= p[0] < 50 and 0 <= p[1] < 50
+    # an assignment uses each row/col at most once
+    assert len({p[0] for p in pairs}) == len(pairs)
+    assert len({p[1] for p in pairs}) == len(pairs)
+
+
+def test_lsap_iou_identity():
+    boxes = np.array(
+        [[0.0, 0.0, 1.0, 1.0], [2.0, 2.0, 3.0, 3.0], [4.0, 4.0, 5.0, 5.0]]
+    )
+    pairs = lsap_iou(boxes, boxes)
+    assert sorted(pairs) == [(0, 0), (1, 1), (2, 2)]
+
+
+def test_lsap_iou_threshold_filters_non_overlap():
+    boxes1 = np.array([[0.0, 0.0, 1.0, 1.0]])
+    boxes2 = np.array([[10.0, 10.0, 11.0, 11.0]])
+    # IoU is 0, so any positive threshold should discard the match.
+    assert lsap_iou(boxes1, boxes2, iou_threshold=0.1) == []
+
+
+def test_lsap_iou_bad_inputs():
+    with pytest.raises(TypeError, match=_BOXES_NOT_NP_ARRAY):
+        lsap_iou("foo", "bar")
+
+
+def test_lsap_iou_bad_dtype():
+    boxes1 = np.random.random((10, 4)).astype(unsuported_dtype_example)
+    boxes2 = np.random.random((10, 4)).astype(unsuported_dtype_example)
+    with pytest.raises(TypeError):
+        lsap_iou(boxes1, boxes2)

powerboxesrs/Cargo.toml

@@ -22,6 +22,7 @@ ndarray = ["dep:ndarray"]
 [dependencies]
 ndarray = { version = ">=0.15, <=0.16", features = ["rayon"], optional = true }
 num-traits = "0.2.17"
+perfect-matching = "0.1.0"
 rayon = "1.8.0"
 rstar = "0.11.0"
 wide = "0.7"