From f4cf351309cb29a5f50cc9d880e0563b45932039 Mon Sep 17 00:00:00 2001 From: gcanat Date: Sun, 22 Mar 2026 00:25:10 +0100 Subject: [PATCH 01/19] add parallel_iou_distance_slice --- powerboxesrs/src/iou.rs | 99 +++++++++++++++++++++++------------------ 1 file changed, 56 insertions(+), 43 deletions(-) diff --git a/powerboxesrs/src/iou.rs b/powerboxesrs/src/iou.rs index de0ade2..1bdf069 100644 --- a/powerboxesrs/src/iou.rs +++ b/powerboxesrs/src/iou.rs @@ -6,7 +6,7 @@ use crate::{ utils, }; #[cfg(feature = "ndarray")] -use ndarray::{Array2, ArrayView2, Zip}; +use ndarray::{Array2, ArrayView2}; use num_traits::{Num, ToPrimitive}; use rstar::RTree; @@ -29,7 +29,7 @@ pub fn iou_distance_slice(boxes1: &[N], boxes2: &[N], n1: usize, n2: usize) - where N: Num + PartialOrd + ToPrimitive + Copy, { - let mut result = vec![0.0f64; n1 * n2]; + let mut result = vec![utils::ONE; n1 * n2]; let areas1 = boxes::box_areas_slice(boxes1, n1); let areas2 = boxes::box_areas_slice(boxes2, n2); @@ -45,7 +45,6 @@ where let x2 = utils::min(a1_x2, a2_x2); let y2 = utils::min(a1_y2, a2_y2); if x2 < x1 || y2 < y1 { - result[i * n2 + j] = utils::ONE; continue; } let intersection = (x2 - x1) * (y2 - y1); @@ -58,6 +57,52 @@ where result } +/// Calculates the intersection over union (IoU) distance between two sets of bounding boxes, +/// in parallel using Rayon. +/// +/// # Arguments +/// +/// * `boxes1` - A flat slice of length `n1 * 4` representing N bounding boxes in xyxy format (row-major). +/// * `boxes2` - A flat slice of length `n2 * 4` representing M bounding boxes in xyxy format (row-major). +/// * `n1` - The number of boxes in the first set. +/// * `n2` - The number of boxes in the second set. +/// +/// # Returns +/// +/// A flat `Vec` of length `n1 * n2` (row-major) representing the IoU distance +/// between each pair of bounding boxes. +pub fn parallel_iou_distance_slice(boxes1: &[N], boxes2: &[N], n1: usize, n2: usize) -> Vec +where + N: Num + PartialOrd + ToPrimitive + Copy + Sync, +{ + let mut result = vec![utils::ONE; n1 * n2]; + let areas1 = boxes::box_areas_slice(boxes1, n1); + let areas2 = boxes::box_areas_slice(boxes2, n2); + + result.par_chunks_mut(n2).enumerate().for_each(|(i, row) | { + let (a1_x1, a1_y1, a1_x2, a1_y2) = utils::row4(boxes1, i); + let area1 = areas1[i]; + + for j in 0..n2 { + let (a2_x1, a2_y1, a2_x2, a2_y2) = utils::row4(boxes2, j); + let area2 = areas2[j]; + let x1 = utils::max(a1_x1, a2_x1); + let y1 = utils::max(a1_y1, a2_y1); + let x2 = utils::min(a1_x2, a2_x2); + let y2 = utils::min(a1_y2, a2_y2); + if x2 < x1 || y2 < y1 { + continue; + } + let intersection = (x2 - x1) * (y2 - y1); + let intersection = intersection.to_f64().unwrap(); + let intersection = utils::min(intersection, utils::min(area1, area2)); + row[j] = utils::ONE - (intersection / (area1 + area2 - intersection)); + } + }); + + result +} + /// Calculates the Rotated Intersection over Union (IoU) distance between two sets of rotated bounding boxes. /// /// # Arguments @@ -274,46 +319,14 @@ where N: Num + PartialEq + PartialOrd + ToPrimitive + Send + Sync + Copy + 'a, BA: Into>, { - let boxes1 = boxes1.into(); - let boxes2 = boxes2.into(); - let num_boxes1 = boxes1.nrows(); - let num_boxes2 = boxes2.nrows(); - - let mut iou_matrix = Array2::::zeros((num_boxes1, num_boxes2)); - let areas_boxes1 = boxes::box_areas(boxes1); - let areas_boxes2 = boxes::box_areas(boxes2); - Zip::indexed(iou_matrix.rows_mut()).par_for_each(|i, mut row| { - let a1 = boxes1.row(i); - let a1_x1 = a1[0]; - let a1_y1 = a1[1]; - let a1_x2 = a1[2]; - let a1_y2 = a1[3]; - let area1 = areas_boxes1[i]; - row.indexed_iter_mut() - .zip(boxes2.rows()) - .for_each(|((j, d), box2)| { - let a2_x1 = box2[0]; - let a2_y1 = box2[1]; - let a2_x2 = box2[2]; - let a2_y2 = box2[3]; - let area2 = areas_boxes2[j]; - - let x1 = utils::max(a1_x1, a2_x1); - let y1 = utils::max(a1_y1, a2_y1); - let x2 = utils::min(a1_x2, a2_x2); - let y2 = utils::min(a1_y2, a2_y2); - if x2 < x1 || y2 < y1 { - *d = utils::ONE; - } else { - let intersection = (x2 - x1) * (y2 - y1); - let intersection = intersection.to_f64().unwrap(); - let intersection = utils::min(intersection, utils::min(area1, area2)); - *d = 1. - (intersection / (area1 + area2 - intersection)); - } - }); - }); - - iou_matrix + let b1 = boxes1.into(); + let b2 = boxes2.into(); + let n1 = b1.nrows(); + let n2 = b2.nrows(); + let s1 = b1.as_slice().expect("boxes1 must be contiguous"); + let s2 = b2.as_slice().expect("boxes2 must be contiguous"); + let result = parallel_iou_distance_slice(s1, s2, n1, n2); + Array2::from_shape_vec((n1, n2), result).unwrap() } /// Calculates the Rotated Intersection over Union (IoU) distance matrix between two sets of rotated bounding boxes. From 096625f59c41c1b6090c240dc45f5782db1e10e7 Mon Sep 17 00:00:00 2001 From: gcanat Date: Sun, 22 Mar 2026 19:36:34 +0100 Subject: [PATCH 02/19] add hungarian_matching --- powerboxesrs/src/assignments.rs | 340 ++++++++++++++++++++++++++++++++ powerboxesrs/src/lib.rs | 1 + 2 files changed, 341 insertions(+) create mode 100644 powerboxesrs/src/assignments.rs diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs new file mode 100644 index 0000000..39e36b2 --- /dev/null +++ b/powerboxesrs/src/assignments.rs @@ -0,0 +1,340 @@ +use crate::iou::{iou_distance_slice, parallel_iou_distance_slice}; +#[cfg(feature = "ndarray")] +use ndarray::ArrayView2; +use num_traits::{Num, ToPrimitive}; + +const SCALE: f64 = 1_000_000_000_000_000.0; + +/// Checks if b < a. Sets a = min(a, b). Returns true if b < a. +#[inline] +fn ckmin(a: &mut T, b: T) -> bool { + if b < *a { + *a = b; + true + } else { + false + } +} + +/// Performs the Hungarian matching algorithm. +/// Adapted from the C++ implementation at https://en.wikipedia.org/wiki/Hungarian_algorithm +/// +/// Given `J` jobs and `W` workers (`J <= W`), computes the minimum cost to assign each jobs +/// to distinct workers. +/// +/// # Arguments +/// * `weights` - A slice representing a `J x W` cost matrix where `c[j][w]` is the cost to +/// assign job `j` to worker `w`. The slice is a row-major representation of the cost matrix. +/// +/// # Returns +/// A `Vec` of length `J`, where entry `j` is the worker's index assigned to this job. +/// +/// # Panics +/// Panics if `weights` is empty, rows have inconsistent lengths, or `J > W`. +/// +/// # Examples +/// +/// ``` +/// use powerboxesrs::assignments::hungarian_matching; +/// let costs = vec![8_i64, 5, 9, 4, 2, 4, 7, 3, 8]; +/// let assignments = hungarian_matching(&costs, 3, 3); +/// assert_eq!(assignments, vec![0, 2, 1]); +/// ``` +pub fn hungarian_matching(weights: &[T], n_rows: usize, n_cols: usize) -> Vec +where + T: Copy + Ord + Default + std::ops::Add + std::ops::Sub, + T: num_traits::Bounded, +{ + assert!( + n_rows <= n_cols, + "Number of jobs must not exceed number of workers" + ); + + // job[w] = job assigned to w-th worker, or None if unassigned. + // A virtual (W+1)-th worker slot is appended for convenience. + let mut job: Vec> = vec![None; n_cols + 1]; + // job potentials + let mut ys: Vec = vec![T::default(); n_rows]; + // worker potentials: -yt[n_cols] will accumulate the sum of all deltas. + let mut yt: Vec = vec![T::default(); n_cols + 1]; + + let inf = T::max_value(); + + for j_cur in 0..n_rows { + // Assign j_cur-th job by routing through the virtual worker slot W. + let mut w_cur = n_cols; + job[w_cur] = Some(j_cur); + + // min reduced cost over edges from the augmenting-path set Z to each worker + let mut min_to: Vec = vec![inf; n_cols + 1]; + // previous worker on the alternating path + let mut prev: Vec> = vec![None; n_cols + 1]; + // whether each worker is currently in the set Z + let mut in_z: Vec = vec![false; n_cols + 1]; + + // Augment: runs at most j_cur + 1 times + while job[w_cur].is_some() { + in_z[w_cur] = true; + let j = job[w_cur].unwrap(); + let mut delta = inf; + let mut w_next = 0usize; + + for w in 0..n_cols { + if !in_z[w] { + let reduced = weights[j * n_cols + w] - ys[j] - yt[w]; + if ckmin(&mut min_to[w], reduced) { + prev[w] = Some(w_cur); + } + if ckmin(&mut delta, min_to[w]) { + w_next = w; + } + } + } + + // Update potentials for all workers + for w in 0..=n_cols { + if in_z[w] { + if let Some(jw) = job[w] { + ys[jw] = ys[jw] + delta; + } + yt[w] = yt[w] - delta; + } else { + min_to[w] = min_to[w] - delta; + } + } + + w_cur = w_next; + } + + // Update job assignments along the alternating path back to W + while w_cur != n_cols { + let w_prev = prev[w_cur].unwrap(); + job[w_cur] = job[w_prev]; + w_cur = w_prev; + } + } + job.pop(); + job.iter().flat_map(|x| *x).collect::>() +} + +/// Compute the optimal assignment between two sets of axis-aligned bounding boxes +/// using the Hungarian algorithm (Kuhn-Munkres), minimising the total IoU distance. +/// +/// Given `n1` ground-truth boxes and `n2` predicted boxes the function builds the +/// `min(n1,n2) × max(n1,n2)` cost matrix from `iou_distance_slice`, scales the +/// `f64` costs to `i64` (multiplied by `1e9`), and calls `kuhn_munkres_min`. +/// After matching, pairs whose IoU is strictly below `iou_threshold` are discarded. +/// +/// # Arguments +/// +/// * `boxes1` - Flat slice of length `n1 * 4` (xyxy, row-major). +/// * `boxes2` - Flat slice of length `n2 * 4` (xyxy, row-major). +/// * `n1` - Number of boxes in the first set. +/// * `n2` - Number of boxes in the second set. +/// * `iou_threshold` - Minimum IoU required to keep a match. Use `0.0` to keep all. +/// +/// # Returns +/// +/// A pair `(indices1, indices2)` of equal-length `Vec` such that +/// `boxes1[indices1[k]]` is matched to `boxes2[indices2[k]]`. +/// The length of both vectors is at most `min(n1, n2)`. +pub fn hungarian_matching_iou_slice( + boxes1: &[N], + boxes2: &[N], + n1: usize, + n2: usize, + iou_threshold: f64, +) -> Vec<(usize, usize)> +where + N: Num + PartialOrd + ToPrimitive + Copy + Sync, +{ + if n1 == 0 || n2 == 0 { + return vec![]; + } + + // benchmark showed that parallel iou distance can be faster above 300 x 300 boxes + let iou_func = if n1 * n2 > 90_000 { + parallel_iou_distance_slice + } else { + iou_distance_slice + }; + + // kuhn_munkres_min requires rows <= columns; transpose when n1 > n2. + let transposed = n1 > n2; + let (nrows, ncols) = if transposed { (n2, n1) } else { (n1, n2) }; + + // Build the cost matrix + let iou_dist = if transposed { + iou_func(boxes1, boxes2, n1, n2) + } else { + iou_func(boxes2, boxes1, n2, n1) + }; + // kuhn_munkres_min expects i64 type, so we scale the values before type casting + let costs_flat: Vec = iou_dist.iter().map(|&d| (d * SCALE) as i64).collect(); + + // let matrix = Matrix::from_vec(nrows, ncols, costs_flat).unwrap(); + + // let (_, assignments) = kuhn_munkres_min(&matrix); + let assignments = hungarian_matching(&costs_flat, nrows, ncols); + + let (raw_idx1, raw_idx2) = if transposed { + (assignments, (0..nrows).collect()) + } else { + ((0..nrows).collect(), assignments) + }; + + // Discard pairs whose IoU falls below the threshold. + let max_dist = 1.0 - iou_threshold; + raw_idx1 + .into_iter() + .zip(raw_idx2) + .filter(|&(i, j)| iou_dist[i * ncols + j] <= max_dist) + .collect() +} + +/// Compute the optimal assignment between two sets of axis-aligned bounding boxes +/// using the Hungarian algorithm, minimising the total IoU distance. +/// +/// Wraps [`hungarian_matching_iou_slice`] for `ndarray` inputs. +/// +/// # Arguments +/// +/// * `boxes1` - Array of shape `(N, 4)` in xyxy format. +/// * `boxes2` - Array of shape `(M, 4)` in xyxy format. +/// * `iou_threshold` - Minimum IoU required to keep a match. Use `0.0` to keep all. +/// +/// # Returns +/// +/// A pair `(indices1, indices2)` of equal-length `Vec` of length at most `min(N, M)`. +#[cfg(feature = "ndarray")] +pub fn hungarian_matching_iou<'a, N, BA>( + boxes1: BA, + boxes2: BA, + iou_threshold: f64, +) -> Vec<(usize, usize)> +where + N: Num + PartialOrd + ToPrimitive + Copy + Sync + 'a, + BA: Into>, +{ + let b1 = boxes1.into(); + let b2 = boxes2.into(); + let n1 = b1.nrows(); + let n2 = b2.nrows(); + let s1 = b1.as_slice().expect("boxes1 must be contiguous"); + let s2 = b2.as_slice().expect("boxes2 must be contiguous"); + hungarian_matching_iou_slice(s1, s2, n1, n2, iou_threshold) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_identical_boxes() { + let boxes = vec![0.0_f64, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0]; + let pairs = hungarian_matching_iou_slice(&boxes, &boxes, 2, 2, 0.0); + assert_eq!(pairs.len(), 2); + for pair in pairs.iter() { + assert_eq!(pair.0, pair.1); + } + } + + #[test] + fn test_more_gt_than_pred() { + let gt = vec![ + 0.0_f64, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0, 5.0, 5.0, + ]; + let pred = vec![ + 2.0, 2.0, 3.0, 3.0, // match with gt 1 + 0.0_f64, 0.0, 1.0, 1.0, // match with gt 0 + ]; + let pairs = hungarian_matching_iou_slice(>, &pred, 3, 2, 0.0); + assert_eq!(pairs.len(), 2); + assert_eq!(pairs, vec![(1, 0), (0, 1)]); + } + + #[test] + fn test_more_pred_than_gt() { + let gt = vec![0.0_f64, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0]; + let pred = vec![ + 0.0_f64, 0.0, 1.0, 1.0, // match with gt 0 + 2.0, 2.0, 3.0, 3.0, // match with gt 1 + 4.0, 4.0, 5.0, 5.0, // no match + ]; + let pairs = hungarian_matching_iou_slice(>, &pred, 2, 3, 0.0); + assert_eq!(pairs.len(), 2); + assert_eq!(pairs, vec![(0, 0), (1, 1)]); + } + + #[test] + fn test_multiple_overlap() { + let gt = vec![ + 0.0_f64, 0.0, 4.0, 4.0, 2.0, 2.0, 4.0, 4.0, 3.0, 3.0, 4.0, 4.0, + ]; + let pred = vec![ + 1.0_f64, 1.0, 3.0, 3.0, // match with gt 1 + 1.0, 1.0, 4.0, 4.0, // match with gt 0 + 2.5, 2.5, 4.0, 4.0, // match with gt 2 + ]; + let pairs = hungarian_matching_iou_slice(>, &pred, 3, 3, 0.0); + assert_eq!(pairs.len(), 3); + assert_eq!(pairs, vec![(0, 1), (1, 0), (2, 2)]); + } + + #[test] + fn test_empty_inputs() { + let boxes: Vec = vec![]; + let pairs = hungarian_matching_iou_slice::(&boxes, &boxes, 0, 0, 0.0); + assert!(pairs.is_empty()); + } + + #[test] + fn test_single_pair() { + let b1 = vec![0.0_f64, 0.0, 2.0, 2.0]; + let b2 = vec![1.0_f64, 1.0, 3.0, 3.0]; + let pairs = hungarian_matching_iou_slice(&b1, &b2, 1, 1, 0.0); + assert_eq!(pairs, vec![(0, 0)]); + } + + #[test] + fn test_optimal_over_greedy() { + // Arrange boxes so that the greedy (nearest-first) choice is suboptimal. + // + // gt0 = [0,0,2,2], gt1 = [3,3,5,5] + // p0 = [3,3,5,5], p1 = [0,0,2,2] + // + // Optimal: gt0→p1, gt1→p0 (total cost 0). Greedy starting at gt0 might pick p0. + let gt = vec![0.0_f64, 0.0, 2.0, 2.0, 3.0, 3.0, 5.0, 5.0]; + let pred = vec![3.0_f64, 3.0, 5.0, 5.0, 0.0, 0.0, 2.0, 2.0]; + let pairs = hungarian_matching_iou_slice(>, &pred, 2, 2, 0.0); + // gt0 should match pred1, gt1 should match pred0. + assert_eq!(pairs, vec![(0, 1), (1, 0)]); + } + + #[test] + fn test_no_overlap_no_match() { + let gt = vec![0.0_f64, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0]; + let pred = vec![10.0_f64, 10.0, 11.0, 11.0, 12.0, 12.0, 13.0, 13.0]; + // With a non-zero IoU threshold, every candidate match should be discarded. + let pairs = hungarian_matching_iou_slice(>, &pred, 2, 2, 0.1); + assert!( + pairs.is_empty(), + "expected no matches when IoU is 0 for all pairs, got {:?}", + pairs + ); + } + + #[cfg(feature = "ndarray")] + mod ndarray_tests { + use super::*; + use ndarray::arr2; + + #[test] + fn test_ndarray_wrapper() { + let boxes1 = arr2(&[[0.0_f64, 0.0, 1.0, 1.0], [2.0, 2.0, 3.0, 3.0]]); + let boxes2 = arr2(&[[0.0_f64, 0.0, 1.0, 1.0], [2.0, 2.0, 3.0, 3.0]]); + let pairs = hungarian_matching_iou(&boxes1, &boxes2, 0.0); + assert_eq!(pairs, vec![(0, 0), (1, 1)]); + } + } +} diff --git a/powerboxesrs/src/lib.rs b/powerboxesrs/src/lib.rs index a7f7b4b..36501cb 100644 --- a/powerboxesrs/src/lib.rs +++ b/powerboxesrs/src/lib.rs @@ -72,6 +72,7 @@ //! - `draw_boxes`: Draw axis-aligned bounding boxes on a CHW image tensor //! - `draw_rotated_boxes`: Draw rotated bounding boxes on a CHW image tensor //! +pub mod assignments; pub mod boxes; pub mod ciou; pub mod diou; From e2c8c4844f08a6e28cf1988c529c0f62365c9a86 Mon Sep 17 00:00:00 2001 From: gcanat Date: Sun, 22 Mar 2026 22:11:35 +0100 Subject: [PATCH 03/19] hungarian_matching_iou python bindings --- bindings/Cargo.lock | 10 +++---- bindings/python/powerboxes/__init__.py | 21 +++++++++++++ bindings/python/powerboxes/_dispatch.py | 1 + bindings/src/lib.rs | 40 ++++++++++++++++++++++++- 4 files changed, 66 insertions(+), 6 deletions(-) diff --git a/bindings/Cargo.lock b/bindings/Cargo.lock index f04f1d3..c9028fa 100644 --- a/bindings/Cargo.lock +++ b/bindings/Cargo.lock @@ -202,9 +202,9 @@ dependencies = [ [[package]] name = "num-traits" -version = "0.2.17" +version = "0.2.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", "libm", @@ -246,7 +246,7 @@ checksum = "f59e70c4aef1e55797c2e8fd94a4f2a973fc972cfde0e0b05f683667b0cd39dd" [[package]] name = "powerboxesrs" -version = "0.3.0" +version = "0.3.1" dependencies = [ "ndarray", "num-traits", @@ -403,9 +403,9 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "semver" -version = "1.0.20" +version = "1.0.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "836fa6a3e1e547f9a2c4040802ec865b5d85f4014efe00555d7090a3dcaa1090" +checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" [[package]] name = "smallvec" diff --git a/bindings/python/powerboxes/__init__.py b/bindings/python/powerboxes/__init__.py index ba3d89e..1b9a053 100644 --- a/bindings/python/powerboxes/__init__.py +++ b/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_hungarian_matching_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 hungarian_matching_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. + + 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_hungarian_matching_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", + "hungarian_matching_iou", "supported_dtypes", "__version__", ] diff --git a/bindings/python/powerboxes/_dispatch.py b/bindings/python/powerboxes/_dispatch.py index 5bb161c..c8da2b2 100644 --- a/bindings/python/powerboxes/_dispatch.py +++ b/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_hungarian_matching_iou = _build_dispatch("hungarian_matching_iou") diff --git a/bindings/src/lib.rs b/bindings/src/lib.rs index 043fed4..b5ac6a0 100644 --- a/bindings/src/lib.rs +++ b/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::{boxes, ciou, diou, draw, giou, iou, nms, tiou, assignments}; 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) -> Array2`. +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> { + $generic(boxes1, boxes2, iou_threshold) + } + } + }; +} + /// Generate a typed `#[pyfunction]` for `(py, boxes) -> Array1`. macro_rules! impl_unary_f64_fn { ($prefix:ident, $generic:ident, $T:ty, $suffix:ident) => { @@ -208,6 +224,12 @@ 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]); + // Hungarian Matching on IoU + register_typed!( + m, + hungarian_matching_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 +820,19 @@ for_each_numeric_type!( rtree_rotated_nms_generic, signed ); + +// Hungarian Matching on IoU +fn hungarian_matching_iou_generic( + boxes1: &Bound<'_, PyArray2>, + boxes2: &Bound<'_, PyArray2>, + iou_threshold: f64, +) -> PyResult> +where + T: Num + ToPrimitive + PartialOrd + numpy::Element + Copy, +{ + let boxes1 = preprocess_boxes(boxes1).unwrap(); + let boxes2 = preprocess_boxes(boxes2).unwrap(); + let asgmt = assignments::hungarian_matching_iou(boxes1, boxes2, iou_threshold); + Ok(asgmt) +} +for_each_numeric_type!(impl_assignment_fn, hungarian_matching_iou, hungarian_matching_iou_generic); From 98cfef59dcf97f56a9ccb5210c9499a788fffb3e Mon Sep 17 00:00:00 2001 From: gcanat Date: Sun, 22 Mar 2026 23:24:41 +0100 Subject: [PATCH 04/19] remove dead code --- powerboxesrs/src/assignments.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index 39e36b2..2825891 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -169,12 +169,9 @@ where } else { iou_func(boxes2, boxes1, n2, n1) }; - // kuhn_munkres_min expects i64 type, so we scale the values before type casting + // hungarian_matching function needs a numeric type that implements the `Ord` trait let costs_flat: Vec = iou_dist.iter().map(|&d| (d * SCALE) as i64).collect(); - // let matrix = Matrix::from_vec(nrows, ncols, costs_flat).unwrap(); - - // let (_, assignments) = kuhn_munkres_min(&matrix); let assignments = hungarian_matching(&costs_flat, nrows, ncols); let (raw_idx1, raw_idx2) = if transposed { From 7db3735e3d8bb069804d2c9d536f98de1cb13ea0 Mon Sep 17 00:00:00 2001 From: gcanat Date: Tue, 24 Mar 2026 00:08:40 +0100 Subject: [PATCH 05/19] replace kuhn_munkres with lsap --- bindings/python/powerboxes/__init__.py | 12 +- bindings/python/powerboxes/_dispatch.py | 2 +- bindings/src/lib.rs | 10 +- powerboxesrs/src/assignments.rs | 195 ++++++++++++------------ 4 files changed, 106 insertions(+), 113 deletions(-) diff --git a/bindings/python/powerboxes/__init__.py b/bindings/python/powerboxes/__init__.py index 1b9a053..51a0b86 100644 --- a/bindings/python/powerboxes/__init__.py +++ b/bindings/python/powerboxes/__init__.py @@ -18,7 +18,7 @@ _dtype_to_func_rtree_nms, _dtype_to_func_rtree_rotated_nms, _dtype_to_func_tiou_distance, - _dtype_to_func_hungarian_matching_iou, + _dtype_to_func_lsap_iou, ) from ._powerboxes import draw_boxes as _draw_boxes from ._powerboxes import draw_rotated_boxes as _draw_rotated_boxes @@ -647,11 +647,13 @@ def draw_rotated_boxes( return _draw_rotated_boxes(image, boxes, colors, thickness, filled, opacity) -def hungarian_matching_iou( +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) @@ -661,9 +663,7 @@ def hungarian_matching_iou( 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_hungarian_matching_iou, boxes1, boxes2, iou_threshold - ) + return _dispatch2(_dtype_to_func_lsap_iou, boxes1, boxes2, iou_threshold) __all__ = [ @@ -690,7 +690,7 @@ def hungarian_matching_iou( "rtree_rotated_nms", "draw_boxes", "draw_rotated_boxes", - "hungarian_matching_iou", + "lsap_iou", "supported_dtypes", "__version__", ] diff --git a/bindings/python/powerboxes/_dispatch.py b/bindings/python/powerboxes/_dispatch.py index c8da2b2..0e1d490 100644 --- a/bindings/python/powerboxes/_dispatch.py +++ b/bindings/python/powerboxes/_dispatch.py @@ -40,4 +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_hungarian_matching_iou = _build_dispatch("hungarian_matching_iou") +_dtype_to_func_lsap_iou = _build_dispatch("lsap_iou") diff --git a/bindings/src/lib.rs b/bindings/src/lib.rs index b5ac6a0..be70e3f 100644 --- a/bindings/src/lib.rs +++ b/bindings/src/lib.rs @@ -227,7 +227,7 @@ fn _powerboxes(m: &Bound<'_, PyModule>) -> PyResult<()> { // Hungarian Matching on IoU register_typed!( m, - hungarian_matching_iou, + lsap_iou, [f64, f32, i64, i32, i16, u64, u32, u16, u8] ); // Masks to boxes @@ -821,8 +821,8 @@ for_each_numeric_type!( signed ); -// Hungarian Matching on IoU -fn hungarian_matching_iou_generic( +// Linear Sum Assignments IoU +fn lsap_iou_generic( boxes1: &Bound<'_, PyArray2>, boxes2: &Bound<'_, PyArray2>, iou_threshold: f64, @@ -832,7 +832,7 @@ where { let boxes1 = preprocess_boxes(boxes1).unwrap(); let boxes2 = preprocess_boxes(boxes2).unwrap(); - let asgmt = assignments::hungarian_matching_iou(boxes1, boxes2, iou_threshold); + let asgmt = assignments::lsap_iou(boxes1, boxes2, iou_threshold); Ok(asgmt) } -for_each_numeric_type!(impl_assignment_fn, hungarian_matching_iou, hungarian_matching_iou_generic); +for_each_numeric_type!(impl_assignment_fn, lsap_iou, lsap_iou_generic); diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index 2825891..2dfb74f 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -5,25 +5,15 @@ use num_traits::{Num, ToPrimitive}; const SCALE: f64 = 1_000_000_000_000_000.0; -/// Checks if b < a. Sets a = min(a, b). Returns true if b < a. -#[inline] -fn ckmin(a: &mut T, b: T) -> bool { - if b < *a { - *a = b; - true - } else { - false - } -} - -/// Performs the Hungarian matching algorithm. -/// Adapted from the C++ implementation at https://en.wikipedia.org/wiki/Hungarian_algorithm +/// Shortest Augmenting Path algorithm for the Linear Sum Assignment Problem. /// +/// Based on: Crouse, "On implementing 2D rectangular assignment algorithms", +/// https://ui.adsabs.harvard.edu/abs/2016ITAES..52.1679C/abstract /// Given `J` jobs and `W` workers (`J <= W`), computes the minimum cost to assign each jobs /// to distinct workers. /// /// # Arguments -/// * `weights` - A slice representing a `J x W` cost matrix where `c[j][w]` is the cost to +/// * `c` - A slice representing a `J x W` cost matrix where `c[j][w]` is the cost to /// assign job `j` to worker `w`. The slice is a row-major representation of the cost matrix. /// /// # Returns @@ -35,94 +25,101 @@ fn ckmin(a: &mut T, b: T) -> bool { /// # Examples /// /// ``` -/// use powerboxesrs::assignments::hungarian_matching; +/// use powerboxesrs::assignments::lsap; /// let costs = vec![8_i64, 5, 9, 4, 2, 4, 7, 3, 8]; -/// let assignments = hungarian_matching(&costs, 3, 3); +/// let assignments = lsap(&costs, 3, 3); /// assert_eq!(assignments, vec![0, 2, 1]); /// ``` -pub fn hungarian_matching(weights: &[T], n_rows: usize, n_cols: usize) -> Vec +pub fn lsap(c: &[T], nrow: usize, ncol: usize) -> Vec where - T: Copy + Ord + Default + std::ops::Add + std::ops::Sub, - T: num_traits::Bounded, + T: Copy + PartialOrd + std::ops::Sub + std::ops::Add, + T: num_traits::Bounded + num_traits::Zero, { - assert!( - n_rows <= n_cols, - "Number of jobs must not exceed number of workers" - ); - - // job[w] = job assigned to w-th worker, or None if unassigned. - // A virtual (W+1)-th worker slot is appended for convenience. - let mut job: Vec> = vec![None; n_cols + 1]; - // job potentials - let mut ys: Vec = vec![T::default(); n_rows]; - // worker potentials: -yt[n_cols] will accumulate the sum of all deltas. - let mut yt: Vec = vec![T::default(); n_cols + 1]; + assert!(nrow <= ncol); let inf = T::max_value(); - for j_cur in 0..n_rows { - // Assign j_cur-th job by routing through the virtual worker slot W. - let mut w_cur = n_cols; - job[w_cur] = Some(j_cur); - - // min reduced cost over edges from the augmenting-path set Z to each worker - let mut min_to: Vec = vec![inf; n_cols + 1]; - // previous worker on the alternating path - let mut prev: Vec> = vec![None; n_cols + 1]; - // whether each worker is currently in the set Z - let mut in_z: Vec = vec![false; n_cols + 1]; - - // Augment: runs at most j_cur + 1 times - while job[w_cur].is_some() { - in_z[w_cur] = true; - let j = job[w_cur].unwrap(); - let mut delta = inf; - let mut w_next = 0usize; - - for w in 0..n_cols { - if !in_z[w] { - let reduced = weights[j * n_cols + w] - ys[j] - yt[w]; - if ckmin(&mut min_to[w], reduced) { - prev[w] = Some(w_cur); + let mut u = vec![T::zero(); nrow]; // row potentials + let mut v = vec![T::zero(); ncol]; // col potentials + let mut col4row = vec![usize::MAX; nrow]; // col assigned to each row + let mut row4col = vec![usize::MAX; ncol]; // row assigned to each col + + for cur_row in 0..nrow { + // Dijkstra-like shortest path from cur_row to any unassigned col + let mut shortest_path_costs = vec![inf; ncol]; + let mut path = vec![usize::MAX; ncol]; + let mut visited = vec![false; ncol]; + + let mut i = cur_row; + let mut sink = usize::MAX; + let mut min_val = T::zero(); + + while sink == usize::MAX { + let mut idx = usize::MAX; + let mut lowest = inf; + + for j in 0..ncol { + if !visited[j] { + let r = c[i * ncol + j] - u[i] - v[j] + min_val; + if r < shortest_path_costs[j] { + shortest_path_costs[j] = r; + path[j] = i; } - if ckmin(&mut delta, min_to[w]) { - w_next = w; + if shortest_path_costs[j] < lowest + || (shortest_path_costs[j] == lowest && row4col[j] == usize::MAX) + { + lowest = shortest_path_costs[j]; + idx = j; } } } - // Update potentials for all workers - for w in 0..=n_cols { - if in_z[w] { - if let Some(jw) = job[w] { - ys[jw] = ys[jw] + delta; - } - yt[w] = yt[w] - delta; - } else { - min_to[w] = min_to[w] - delta; - } + min_val = lowest; + let j = idx; + visited[j] = true; + + if row4col[j] == usize::MAX { + sink = j; + } else { + i = row4col[j]; } + } - w_cur = w_next; + // Update potentials along the path + u[cur_row] = u[cur_row] + min_val; + for j in 0..ncol { + if visited[j] { + let r = row4col[j]; + if r != usize::MAX { + u[r] = u[r] - shortest_path_costs[j] + min_val; + } + v[j] = v[j] + shortest_path_costs[j] - min_val; + } } - // Update job assignments along the alternating path back to W - while w_cur != n_cols { - let w_prev = prev[w_cur].unwrap(); - job[w_cur] = job[w_prev]; - w_cur = w_prev; + // Augment along the path back to cur_row + let mut j = sink; + loop { + let i = path[j]; + row4col[j] = i; + let prev_j = col4row[i]; + col4row[i] = j; + if i == cur_row { + break; + } + j = prev_j; } } - job.pop(); - job.iter().flat_map(|x| *x).collect::>() + + col4row } /// Compute the optimal assignment between two sets of axis-aligned bounding boxes -/// using the Hungarian algorithm (Kuhn-Munkres), minimising the total IoU distance. +/// using the LSAP algorithm, minimising the total IoU distance. /// /// Given `n1` ground-truth boxes and `n2` predicted boxes the function builds the /// `min(n1,n2) × max(n1,n2)` cost matrix from `iou_distance_slice`, scales the -/// `f64` costs to `i64` (multiplied by `1e9`), and calls `kuhn_munkres_min`. +/// `f64` costs to `i64` (multiplied by `1e15`), and calls `lsap`. /// After matching, pairs whose IoU is strictly below `iou_threshold` are discarded. /// /// # Arguments @@ -138,7 +135,7 @@ where /// A pair `(indices1, indices2)` of equal-length `Vec` such that /// `boxes1[indices1[k]]` is matched to `boxes2[indices2[k]]`. /// The length of both vectors is at most `min(n1, n2)`. -pub fn hungarian_matching_iou_slice( +pub fn lsap_iou_slice( boxes1: &[N], boxes2: &[N], n1: usize, @@ -159,20 +156,20 @@ where iou_distance_slice }; - // kuhn_munkres_min requires rows <= columns; transpose when n1 > n2. + // lsap requires rows <= columns; transpose when n1 > n2. let transposed = n1 > n2; let (nrows, ncols) = if transposed { (n2, n1) } else { (n1, n2) }; // Build the cost matrix - let iou_dist = if transposed { - iou_func(boxes1, boxes2, n1, n2) + let iou_dist = if !transposed { + iou_func(boxes1, boxes2, nrows, ncols) } else { - iou_func(boxes2, boxes1, n2, n1) + iou_func(boxes2, boxes1, nrows, ncols) }; - // hungarian_matching function needs a numeric type that implements the `Ord` trait + // lsap function needs a numeric type that implements the `Ord` trait let costs_flat: Vec = iou_dist.iter().map(|&d| (d * SCALE) as i64).collect(); - let assignments = hungarian_matching(&costs_flat, nrows, ncols); + let assignments = lsap(&costs_flat, nrows, ncols); let (raw_idx1, raw_idx2) = if transposed { (assignments, (0..nrows).collect()) @@ -190,9 +187,9 @@ where } /// Compute the optimal assignment between two sets of axis-aligned bounding boxes -/// using the Hungarian algorithm, minimising the total IoU distance. +/// using the LSAP algorithm, minimising the total IoU distance. /// -/// Wraps [`hungarian_matching_iou_slice`] for `ndarray` inputs. +/// Wraps [`lsap_iou_slice`] for `ndarray` inputs. /// /// # Arguments /// @@ -204,11 +201,7 @@ where /// /// A pair `(indices1, indices2)` of equal-length `Vec` of length at most `min(N, M)`. #[cfg(feature = "ndarray")] -pub fn hungarian_matching_iou<'a, N, BA>( - boxes1: BA, - boxes2: BA, - iou_threshold: f64, -) -> Vec<(usize, usize)> +pub fn lsap_iou<'a, N, BA>(boxes1: BA, boxes2: BA, iou_threshold: f64) -> Vec<(usize, usize)> where N: Num + PartialOrd + ToPrimitive + Copy + Sync + 'a, BA: Into>, @@ -219,7 +212,7 @@ where let n2 = b2.nrows(); let s1 = b1.as_slice().expect("boxes1 must be contiguous"); let s2 = b2.as_slice().expect("boxes2 must be contiguous"); - hungarian_matching_iou_slice(s1, s2, n1, n2, iou_threshold) + lsap_iou_slice(s1, s2, n1, n2, iou_threshold) } #[cfg(test)] @@ -229,7 +222,7 @@ mod tests { #[test] fn test_identical_boxes() { let boxes = vec![0.0_f64, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0]; - let pairs = hungarian_matching_iou_slice(&boxes, &boxes, 2, 2, 0.0); + let pairs = lsap_iou_slice(&boxes, &boxes, 2, 2, 0.0); assert_eq!(pairs.len(), 2); for pair in pairs.iter() { assert_eq!(pair.0, pair.1); @@ -245,7 +238,7 @@ mod tests { 2.0, 2.0, 3.0, 3.0, // match with gt 1 0.0_f64, 0.0, 1.0, 1.0, // match with gt 0 ]; - let pairs = hungarian_matching_iou_slice(>, &pred, 3, 2, 0.0); + let pairs = lsap_iou_slice(>, &pred, 3, 2, 0.0); assert_eq!(pairs.len(), 2); assert_eq!(pairs, vec![(1, 0), (0, 1)]); } @@ -258,7 +251,7 @@ mod tests { 2.0, 2.0, 3.0, 3.0, // match with gt 1 4.0, 4.0, 5.0, 5.0, // no match ]; - let pairs = hungarian_matching_iou_slice(>, &pred, 2, 3, 0.0); + let pairs = lsap_iou_slice(>, &pred, 2, 3, 0.0); assert_eq!(pairs.len(), 2); assert_eq!(pairs, vec![(0, 0), (1, 1)]); } @@ -273,7 +266,7 @@ mod tests { 1.0, 1.0, 4.0, 4.0, // match with gt 0 2.5, 2.5, 4.0, 4.0, // match with gt 2 ]; - let pairs = hungarian_matching_iou_slice(>, &pred, 3, 3, 0.0); + let pairs = lsap_iou_slice(>, &pred, 3, 3, 0.0); assert_eq!(pairs.len(), 3); assert_eq!(pairs, vec![(0, 1), (1, 0), (2, 2)]); } @@ -281,7 +274,7 @@ mod tests { #[test] fn test_empty_inputs() { let boxes: Vec = vec![]; - let pairs = hungarian_matching_iou_slice::(&boxes, &boxes, 0, 0, 0.0); + let pairs = lsap_iou_slice::(&boxes, &boxes, 0, 0, 0.0); assert!(pairs.is_empty()); } @@ -289,7 +282,7 @@ mod tests { fn test_single_pair() { let b1 = vec![0.0_f64, 0.0, 2.0, 2.0]; let b2 = vec![1.0_f64, 1.0, 3.0, 3.0]; - let pairs = hungarian_matching_iou_slice(&b1, &b2, 1, 1, 0.0); + let pairs = lsap_iou_slice(&b1, &b2, 1, 1, 0.0); assert_eq!(pairs, vec![(0, 0)]); } @@ -303,7 +296,7 @@ mod tests { // Optimal: gt0→p1, gt1→p0 (total cost 0). Greedy starting at gt0 might pick p0. let gt = vec![0.0_f64, 0.0, 2.0, 2.0, 3.0, 3.0, 5.0, 5.0]; let pred = vec![3.0_f64, 3.0, 5.0, 5.0, 0.0, 0.0, 2.0, 2.0]; - let pairs = hungarian_matching_iou_slice(>, &pred, 2, 2, 0.0); + let pairs = lsap_iou_slice(>, &pred, 2, 2, 0.0); // gt0 should match pred1, gt1 should match pred0. assert_eq!(pairs, vec![(0, 1), (1, 0)]); } @@ -313,7 +306,7 @@ mod tests { let gt = vec![0.0_f64, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0]; let pred = vec![10.0_f64, 10.0, 11.0, 11.0, 12.0, 12.0, 13.0, 13.0]; // With a non-zero IoU threshold, every candidate match should be discarded. - let pairs = hungarian_matching_iou_slice(>, &pred, 2, 2, 0.1); + let pairs = lsap_iou_slice(>, &pred, 2, 2, 0.1); assert!( pairs.is_empty(), "expected no matches when IoU is 0 for all pairs, got {:?}", @@ -330,7 +323,7 @@ mod tests { fn test_ndarray_wrapper() { let boxes1 = arr2(&[[0.0_f64, 0.0, 1.0, 1.0], [2.0, 2.0, 3.0, 3.0]]); let boxes2 = arr2(&[[0.0_f64, 0.0, 1.0, 1.0], [2.0, 2.0, 3.0, 3.0]]); - let pairs = hungarian_matching_iou(&boxes1, &boxes2, 0.0); + let pairs = lsap_iou(&boxes1, &boxes2, 0.0); assert_eq!(pairs, vec![(0, 0), (1, 1)]); } } From 5cf3f8f2cf6a4340354ad03ff1f8b8a41dfe4c55 Mon Sep 17 00:00:00 2001 From: gcanat Date: Sat, 28 Mar 2026 19:40:36 +0100 Subject: [PATCH 06/19] add simd implementation --- bindings/Cargo.lock | 21 +++ powerboxesrs/Cargo.lock | 21 +++ powerboxesrs/Cargo.toml | 2 + powerboxesrs/src/assignments.rs | 271 +++++++++++++++++++++++++++++++- 4 files changed, 309 insertions(+), 6 deletions(-) diff --git a/bindings/Cargo.lock b/bindings/Cargo.lock index c9028fa..1750f08 100644 --- a/bindings/Cargo.lock +++ b/bindings/Cargo.lock @@ -187,6 +187,7 @@ version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ba157ca0885411de85d6ca030ba7e2a83a28636056c7c699b07c8b6f7383214" dependencies = [ + "bytemuck", "num-traits", ] @@ -248,8 +249,10 @@ checksum = "f59e70c4aef1e55797c2e8fd94a4f2a973fc972cfde0e0b05f683667b0cd39dd" name = "powerboxesrs" version = "0.3.1" dependencies = [ + "bytemuck", "ndarray", "num-traits", + "pulp", "rayon", "rstar", "wide", @@ -264,6 +267,18 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "pulp" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9e3f19bdeda2e49d16c8ae90f9615adc2298ee16974bb250d0afb705e33043f" +dependencies = [ + "bytemuck", + "libm", + "num-complex", + "reborrow", +] + [[package]] name = "pyo3" version = "0.27.2" @@ -360,6 +375,12 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "reborrow" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430" + [[package]] name = "rstar" version = "0.11.0" diff --git a/powerboxesrs/Cargo.lock b/powerboxesrs/Cargo.lock index c59a141..5102d78 100644 --- a/powerboxesrs/Cargo.lock +++ b/powerboxesrs/Cargo.lock @@ -404,6 +404,7 @@ version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ba157ca0885411de85d6ca030ba7e2a83a28636056c7c699b07c8b6f7383214" dependencies = [ + "bytemuck", "num-traits", ] @@ -471,10 +472,12 @@ dependencies = [ name = "powerboxesrs" version = "0.3.1" dependencies = [ + "bytemuck", "codspeed-criterion-compat", "criterion", "ndarray", "num-traits", + "pulp", "rayon", "rstar", "wide", @@ -489,6 +492,18 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "pulp" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9e3f19bdeda2e49d16c8ae90f9615adc2298ee16974bb250d0afb705e33043f" +dependencies = [ + "bytemuck", + "libm", + "num-complex", + "reborrow", +] + [[package]] name = "quote" version = "1.0.33" @@ -524,6 +539,12 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "reborrow" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430" + [[package]] name = "regex" version = "1.10.2" diff --git a/powerboxesrs/Cargo.toml b/powerboxesrs/Cargo.toml index 8b7471b..5024255 100644 --- a/powerboxesrs/Cargo.toml +++ b/powerboxesrs/Cargo.toml @@ -20,8 +20,10 @@ default = ["ndarray"] ndarray = ["dep:ndarray"] [dependencies] +bytemuck = "1.15" ndarray = { version = ">=0.15, <=0.16", features = ["rayon"], optional = true } num-traits = "0.2.17" +pulp = "0.20" rayon = "1.8.0" rstar = "0.11.0" wide = "0.7" diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index 2dfb74f..dce72d2 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -1,9 +1,15 @@ use crate::iou::{iou_distance_slice, parallel_iou_distance_slice}; +use bytemuck::cast_slice; #[cfg(feature = "ndarray")] use ndarray::ArrayView2; use num_traits::{Num, ToPrimitive}; +use pulp::{Arch, Simd, WithSimd}; + +// Cost matrix size above which we use parallel_iou_distance_slice +const PARALLEL_IOU_MIN_BOXES: usize = 90_000; +// If more than 75% boxes dont overlap, lsap_simd is slower than lsap +const SIMD_MAX_SPARSITY: f32 = 0.75; -const SCALE: f64 = 1_000_000_000_000_000.0; /// Shortest Augmenting Path algorithm for the Linear Sum Assignment Problem. /// @@ -114,13 +120,251 @@ where col4row } +/// SIMD implementation for LSAP +struct InnerScan<'a> { + c_row: &'a [f32], + v: &'a [f32], + visited: &'a [u8], + spc: &'a mut [f32], + path: &'a mut [u32], + row4col: &'a [u32], + u_i: f32, + min_val: f32, + row_i: u32, +} + +struct ScanResult { + pub best_cost: f32, + pub best_col: usize, +} + +impl WithSimd for InnerScan<'_> { + type Output = ScanResult; + + #[inline(always)] + fn with_simd(self, simd: S) -> ScanResult { + let Self { + c_row, + v, + visited, + spc, + path, + row4col, + u_i, + min_val, + row_i, + } = self; + let ncol = c_row.len(); + + let offset = simd.splat_f32s(min_val - u_i); + let inf_v = simd.splat_f32s(f32::INFINITY); + let row_v = simd.splat_u32s(row_i); + let zero_u = simd.splat_u32s(0u32); + + let mut best_cost_v = simd.splat_f32s(f32::INFINITY); + let mut best_col_v = simd.splat_u32s(u32::MAX); + + let (c_chunks, c_tail) = S::as_simd_f32s(c_row); + let (v_chunks, _) = S::as_simd_f32s(v); + let (spc_chunks, _) = S::as_mut_simd_f32s(spc); + let (path_chunks, _) = S::as_mut_simd_u32s(path); + + let lanes = std::mem::size_of::() / 4; + + for (chunk_idx, (((c_v, v_v), spc_v), path_v)) in c_chunks + .iter() + .zip(v_chunks.iter()) + .zip(spc_chunks.iter_mut()) + .zip(path_chunks.iter_mut()) + .enumerate() + { + let base = chunk_idx * lanes; + + // Widen u8 visited flags → u32 lanes, then build mask where != 0 + let vis_u32: S::u32s = { + let mut buf = zero_u; + let buf_slice: &mut [u32] = + bytemuck::cast_slice_mut(std::slice::from_mut(&mut buf)); + for (dst, &src) in buf_slice.iter_mut().zip(visited[base..].iter()) { + *dst = src as u32; + } + buf + }; + let is_visited = simd.greater_than_u32s(vis_u32, zero_u); + + // r = c[j] - v[j] + (min_val - u_i) + let r = simd.add_f32s(simd.sub_f32s(*c_v, *v_v), offset); + + // spc[j] = min(spc[j], r) for unvisited lanes only + let old_spc_v = *spc_v; + let new_spc = simd.min_f32s(*spc_v, r); + *spc_v = simd.select_f32s_m32s(is_visited, *spc_v, new_spc); + + // path[j] = row_i where r improved spc AND lane is unvisited + let improved = simd.less_than_f32s(r, old_spc_v); + let update_path = simd.and_m32s(simd.not_m32s(is_visited), improved); + *path_v = simd.select_u32s_m32s(update_path, row_v, *path_v); + + // For argmin: mask visited lanes out with infinity + let cost_for_min = simd.select_f32s_m32s(is_visited, inf_v, *spc_v); + + // Column indices for this chunk + let col_indices: S::u32s = { + let mut buf = zero_u; + let buf_slice: &mut [u32] = + bytemuck::cast_slice_mut(std::slice::from_mut(&mut buf)); + for (k, dst) in buf_slice.iter_mut().enumerate() { + *dst = (base + k) as u32; + } + buf + }; + + // Update running per-lane argmin + let new_is_better = simd.less_than_f32s(cost_for_min, best_cost_v); + best_cost_v = simd.select_f32s_m32s(new_is_better, cost_for_min, best_cost_v); + best_col_v = simd.select_u32s_m32s(new_is_better, col_indices, best_col_v); + } + + // Horizontal reduction: fold SIMD lanes down to a scalar argmin + let costs: &[f32] = cast_slice(std::slice::from_ref(&best_cost_v)); + let cols: &[u32] = cast_slice(std::slice::from_ref(&best_col_v)); + let mut best_cost = f32::INFINITY; + let mut best_col = usize::MAX; + for (&cost, &col) in costs.iter().zip(cols.iter()) { + if cost < best_cost + || (cost == best_cost && col != u32::MAX && row4col[col as usize] == u32::MAX) + { + best_cost = cost; + best_col = col as usize; + } + } + + // Scalar tail: remainder columns that don't fill a full SIMD vector + let tail_start = ncol - c_tail.len(); + for j in tail_start..ncol { + if visited[j] == 0 { + let r = c_row[j] - u_i - v[j] + min_val; + if r < spc[j] { + spc[j] = r; + path[j] = row_i; + } + if spc[j] < best_cost || (spc[j] == best_cost && row4col[j] == u32::MAX) { + best_cost = spc[j]; + best_col = j; + } + } + } + + ScanResult { + best_cost, + best_col, + } + } +} + +/// Shortest Augmenting Path algorithm with SIMD instructions for the Linear Sum Assignment Problem. +/// +/// Based on: Crouse, "On implementing 2D rectangular assignment algorithms", +/// https://ui.adsabs.harvard.edu/abs/2016ITAES..52.1679C/abstract +/// Given `J` jobs and `W` workers (`J <= W`), computes the minimum cost to assign each jobs +/// to distinct workers. +/// +/// # Arguments +/// * `c` - A slice representing a `J x W` cost matrix where `c[j][w]` is the cost to +/// assign job `j` to worker `w`. The slice is a row-major representation of the cost matrix. +/// +/// # Returns +/// A `Vec` of length `J`, where entry `j` is the worker's index assigned to this job. +/// +/// # Panics +/// Panics if `weights` is empty, rows have inconsistent lengths, or `J > W`. +/// +/// # Examples +/// +/// ``` +/// use powerboxesrs::assignments::lsap_simd; +/// let costs = vec![8_f32, 5., 9., 4., 2., 4., 7., 3., 8.]; +/// let assignments = lsap_simd(&costs, 3, 3); +/// assert_eq!(assignments, vec![0, 2, 1]); +/// ``` +pub fn lsap_simd(c: &[f32], nrow: usize, ncol: usize) -> Vec { + assert!(nrow <= ncol); + + let arch = Arch::new(); + + let mut u = vec![0f32; nrow]; + let mut v = vec![0f32; ncol]; + let mut col4row = vec![usize::MAX; nrow]; + let mut row4col = vec![u32::MAX; ncol]; + let mut visited = vec![0u8; ncol]; + + for cur_row in 0..nrow { + let mut spc = vec![f32::INFINITY; ncol]; + let mut path = vec![u32::MAX; ncol]; + visited.fill(0); + + let mut i = cur_row; + let mut sink = usize::MAX; + let mut min_val = 0f32; + + while sink == usize::MAX { + let res = arch.dispatch(InnerScan { + c_row: &c[i * ncol..(i + 1) * ncol], + v: &v, + visited: &visited, + spc: &mut spc, + path: &mut path, + row4col: &row4col, + u_i: u[i], + min_val, + row_i: i as u32, + }); + + min_val = res.best_cost; + let j = res.best_col; + visited[j] = 1; + + if row4col[j] == u32::MAX { + sink = j; + } else { + i = row4col[j] as usize; + } + } + + u[cur_row] += min_val; + for j in 0..ncol { + if visited[j] != 0 { + let r = row4col[j]; + if r != u32::MAX { + u[r as usize] += min_val - spc[j]; + } + v[j] += spc[j] - min_val; + } + } + + let mut j = sink; + loop { + let pi = path[j] as usize; + row4col[j] = pi as u32; + let prev_j = col4row[pi]; + col4row[pi] = j; + if pi == cur_row { + break; + } + j = prev_j; + } + } + + col4row +} + /// Compute the optimal assignment between two sets of axis-aligned bounding boxes /// using the LSAP algorithm, minimising the total IoU distance. /// /// Given `n1` ground-truth boxes and `n2` predicted boxes the function builds the /// `min(n1,n2) × max(n1,n2)` cost matrix from `iou_distance_slice`, scales the /// `f64` costs to `i64` (multiplied by `1e15`), and calls `lsap`. -/// After matching, pairs whose IoU is strictly below `iou_threshold` are discarded. +/// After matching, pairs whose IoU is below `iou_threshold` are discarded. /// /// # Arguments /// @@ -150,7 +394,7 @@ where } // benchmark showed that parallel iou distance can be faster above 300 x 300 boxes - let iou_func = if n1 * n2 > 90_000 { + let iou_func = if n1 * n2 > PARALLEL_IOU_MIN_BOXES { parallel_iou_distance_slice } else { iou_distance_slice @@ -166,10 +410,25 @@ where } else { iou_func(boxes2, boxes1, nrows, ncols) }; - // lsap function needs a numeric type that implements the `Ord` trait - let costs_flat: Vec = iou_dist.iter().map(|&d| (d * SCALE) as i64).collect(); - let assignments = lsap(&costs_flat, nrows, ncols); + let costs_flat: Vec = iou_dist.iter().map(|&d| d as f32).collect(); + + // check matrix "sparsity": ie many boxes dont overlap (iou is 0, distance is 1) + let non_overlapping_ratio = + costs_flat + .iter() + .fold(0_f32, |acc, x| if *x == 1_f32 { &acc + 1_f32 } else { acc }) + / costs_flat.len() as f32; + + // if cost matrix is too sparse, the overhead of calling InnerScan in the SIMD + // implementation dominates so it's much faster to use the scalar implementation. + let lsap_func = if non_overlapping_ratio > SIMD_MAX_SPARSITY { + lsap + } else { + lsap_simd + }; + + let assignments = lsap_func(&costs_flat, nrows, ncols); let (raw_idx1, raw_idx2) = if transposed { (assignments, (0..nrows).collect()) From d7c672856c5a7596df2caad2fce775d133b4d940 Mon Sep 17 00:00:00 2001 From: gcanat Date: Sat, 28 Mar 2026 22:36:12 +0100 Subject: [PATCH 07/19] add simple tests to lsap and lsap_simd --- powerboxesrs/src/assignments.rs | 45 ++++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index dce72d2..442e279 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -10,7 +10,6 @@ const PARALLEL_IOU_MIN_BOXES: usize = 90_000; // If more than 75% boxes dont overlap, lsap_simd is slower than lsap const SIMD_MAX_SPARSITY: f32 = 0.75; - /// Shortest Augmenting Path algorithm for the Linear Sum Assignment Problem. /// /// Based on: Crouse, "On implementing 2D rectangular assignment algorithms", @@ -478,6 +477,49 @@ where mod tests { use super::*; + // tests for lsap and lsap_simd + #[test] + fn test_identical_costs() { + let cost_matrix = vec![1.0_f32, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]; + let asgmt = lsap(&cost_matrix, 3, 3); + let asgmt_simd = lsap_simd(&cost_matrix, 3, 3); + // all costs are equal: matches are made starting from the end + assert_eq!(asgmt, vec![2, 1, 0]); + assert_eq!(asgmt_simd, vec![2, 1, 0]); + } + + #[test] + fn test_known_assignment() { + let costs_f32 = vec![8_f32, 5., 9., 4., 2., 4., 7., 3., 8.]; + assert_eq!(lsap(&costs_f32, 3, 3), vec![0, 2, 1]); + assert_eq!(lsap_simd(&costs_f32, 3, 3), vec![0, 2, 1]); + + let costs_i64 = vec![8_i64, 5, 9, 4, 2, 4, 7, 3, 8]; + assert_eq!(lsap(&costs_i64, 3, 3), vec![0, 2, 1]); + } + + #[test] + fn test_single_element() { + let costs = vec![42.0_f32]; + assert_eq!(lsap(&costs, 1, 1), vec![0]); + assert_eq!(lsap_simd(&costs, 1, 1), vec![0]); + } + + #[test] + fn test_rectangular() { + let costs = vec![3.0_f32, 1., 2., 1., 3., 2.]; + assert_eq!(lsap(&costs, 2, 3), vec![1, 0]); + assert_eq!(lsap_simd(&costs, 2, 3), vec![1, 0]); + } + + #[test] + fn test_optimal_in_diagonal() { + let costs = vec![0.0_f32, 1., 1., 1., 0., 1., 1., 1., 0.]; + assert_eq!(lsap(&costs, 3, 3), vec![0, 1, 2]); + assert_eq!(lsap_simd(&costs, 3, 3), vec![0, 1, 2]); + } + + // tests for lsap_iou_slice #[test] fn test_identical_boxes() { let boxes = vec![0.0_f64, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0]; @@ -573,6 +615,7 @@ mod tests { ); } + // test for ndarray lsap_iou #[cfg(feature = "ndarray")] mod ndarray_tests { use super::*; From b939add46ab9e8eaa899652bed3d927dd74be78c Mon Sep 17 00:00:00 2001 From: gcanat Date: Sun, 29 Mar 2026 00:15:42 +0100 Subject: [PATCH 08/19] small typos --- bindings/src/lib.rs | 4 ++-- powerboxesrs/src/assignments.rs | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/bindings/src/lib.rs b/bindings/src/lib.rs index be70e3f..c2b8cbc 100644 --- a/bindings/src/lib.rs +++ b/bindings/src/lib.rs @@ -68,7 +68,7 @@ macro_rules! impl_distance2_fn { }; } -/// Generate a typed `#[pyfunction]` for `(py, boxes1, boxes2, iou_threshold) -> Array2`. +/// 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! { @@ -224,7 +224,7 @@ 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]); - // Hungarian Matching on IoU + // LSAP on IoU register_typed!( m, lsap_iou, diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index 442e279..f4e7d44 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -7,8 +7,8 @@ use pulp::{Arch, Simd, WithSimd}; // Cost matrix size above which we use parallel_iou_distance_slice const PARALLEL_IOU_MIN_BOXES: usize = 90_000; -// If more than 75% boxes dont overlap, lsap_simd is slower than lsap -const SIMD_MAX_SPARSITY: f32 = 0.75; +// If more than 65% boxes dont overlap, lsap_simd is slower than lsap +const SIMD_MAX_SPARSITY: f32 = 0.65; /// Shortest Augmenting Path algorithm for the Linear Sum Assignment Problem. /// @@ -419,6 +419,8 @@ where .fold(0_f32, |acc, x| if *x == 1_f32 { &acc + 1_f32 } else { acc }) / costs_flat.len() as f32; + println!("Sparsity rust: {}", non_overlapping_ratio); + // if cost matrix is too sparse, the overhead of calling InnerScan in the SIMD // implementation dominates so it's much faster to use the scalar implementation. let lsap_func = if non_overlapping_ratio > SIMD_MAX_SPARSITY { From dd46b93d97b219b416f9808568d0786311758f93 Mon Sep 17 00:00:00 2001 From: gcanat Date: Sun, 29 Mar 2026 00:20:22 +0100 Subject: [PATCH 09/19] formatting --- bindings/src/lib.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/bindings/src/lib.rs b/bindings/src/lib.rs index c2b8cbc..adf7e47 100644 --- a/bindings/src/lib.rs +++ b/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, assignments}; +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}; @@ -225,11 +225,7 @@ fn _powerboxes(m: &Bound<'_, PyModule>) -> PyResult<()> { // 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] - ); + 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 From 7f551d9d90d28eefc6313328e05dff39840245b4 Mon Sep 17 00:00:00 2001 From: gcanat Date: Sun, 29 Mar 2026 00:24:55 +0100 Subject: [PATCH 10/19] more formatting --- powerboxesrs/src/iou.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/powerboxesrs/src/iou.rs b/powerboxesrs/src/iou.rs index 1bdf069..423f2ca 100644 --- a/powerboxesrs/src/iou.rs +++ b/powerboxesrs/src/iou.rs @@ -79,7 +79,7 @@ where let areas1 = boxes::box_areas_slice(boxes1, n1); let areas2 = boxes::box_areas_slice(boxes2, n2); - result.par_chunks_mut(n2).enumerate().for_each(|(i, row) | { + result.par_chunks_mut(n2).enumerate().for_each(|(i, row)| { let (a1_x1, a1_y1, a1_x2, a1_y2) = utils::row4(boxes1, i); let area1 = areas1[i]; From 8447f7d17325cd15075164508fffd232695ae40e Mon Sep 17 00:00:00 2001 From: gcanat Date: Sun, 29 Mar 2026 11:05:55 +0200 Subject: [PATCH 11/19] fix sparsity calculation --- powerboxesrs/src/assignments.rs | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index f4e7d44..b7ce371 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -7,8 +7,8 @@ use pulp::{Arch, Simd, WithSimd}; // Cost matrix size above which we use parallel_iou_distance_slice const PARALLEL_IOU_MIN_BOXES: usize = 90_000; -// If more than 65% boxes dont overlap, lsap_simd is slower than lsap -const SIMD_MAX_SPARSITY: f32 = 0.65; +// If more than 95% boxes dont overlap, lsap_simd is slower than lsap +const SIMD_MAX_SPARSITY: f64 = 0.95; /// Shortest Augmenting Path algorithm for the Linear Sum Assignment Problem. /// @@ -413,13 +413,11 @@ where let costs_flat: Vec = iou_dist.iter().map(|&d| d as f32).collect(); // check matrix "sparsity": ie many boxes dont overlap (iou is 0, distance is 1) - let non_overlapping_ratio = - costs_flat - .iter() - .fold(0_f32, |acc, x| if *x == 1_f32 { &acc + 1_f32 } else { acc }) - / costs_flat.len() as f32; - - println!("Sparsity rust: {}", non_overlapping_ratio); + let non_overlapping_ratio = costs_flat + .iter() + .filter_map(|x| if *x < 1_f32 { None } else { Some(1_f64) }) + .sum::() + / (costs_flat.len() as f64); // if cost matrix is too sparse, the overhead of calling InnerScan in the SIMD // implementation dominates so it's much faster to use the scalar implementation. From 07784e74a0b73cc83f419a9a2582d0c03f02fb5e Mon Sep 17 00:00:00 2001 From: gcanat Date: Tue, 31 Mar 2026 23:30:52 +0200 Subject: [PATCH 12/19] improve simd implementation --- powerboxesrs/src/assignments.rs | 66 ++++++++++++++++----------------- 1 file changed, 31 insertions(+), 35 deletions(-) diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index b7ce371..e2fb7e7 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -7,8 +7,8 @@ use pulp::{Arch, Simd, WithSimd}; // Cost matrix size above which we use parallel_iou_distance_slice const PARALLEL_IOU_MIN_BOXES: usize = 90_000; -// If more than 95% boxes dont overlap, lsap_simd is slower than lsap -const SIMD_MAX_SPARSITY: f64 = 0.95; +// If more than 98% boxes dont overlap, lsap_simd is slower than lsap +const SIMD_MAX_SPARSITY: f64 = 0.98; /// Shortest Augmenting Path algorithm for the Linear Sum Assignment Problem. /// @@ -80,13 +80,12 @@ where } min_val = lowest; - let j = idx; - visited[j] = true; + visited[idx] = true; - if row4col[j] == usize::MAX { - sink = j; + if row4col[idx] == usize::MAX { + sink = idx; } else { - i = row4col[j]; + i = row4col[idx]; } } @@ -118,12 +117,11 @@ where col4row } - /// SIMD implementation for LSAP struct InnerScan<'a> { c_row: &'a [f32], v: &'a [f32], - visited: &'a [u8], + visited: &'a [u32], spc: &'a mut [f32], path: &'a mut [u32], row4col: &'a [u32], @@ -167,29 +165,33 @@ impl WithSimd for InnerScan<'_> { let (v_chunks, _) = S::as_simd_f32s(v); let (spc_chunks, _) = S::as_mut_simd_f32s(spc); let (path_chunks, _) = S::as_mut_simd_u32s(path); + let (vis_chunks, _) = S::as_simd_u32s(visited); let lanes = std::mem::size_of::() / 4; - for (chunk_idx, (((c_v, v_v), spc_v), path_v)) in c_chunks + // Precompute iota = [0, 1, 2, ..., lanes-1] once. + // Per-chunk indices are derived as iota + splat(base) + let iota: S::u32s = { + let mut buf = zero_u; + let buf_slice: &mut [u32] = bytemuck::cast_slice_mut(std::slice::from_mut(&mut buf)); + for (k, dst) in buf_slice.iter_mut().enumerate() { + *dst = k as u32; + } + buf + }; + + for (chunk_idx, ((((c_v, v_v), spc_v), path_v), vis_v)) in c_chunks .iter() .zip(v_chunks.iter()) .zip(spc_chunks.iter_mut()) .zip(path_chunks.iter_mut()) + .zip(vis_chunks.iter()) .enumerate() { let base = chunk_idx * lanes; - // Widen u8 visited flags → u32 lanes, then build mask where != 0 - let vis_u32: S::u32s = { - let mut buf = zero_u; - let buf_slice: &mut [u32] = - bytemuck::cast_slice_mut(std::slice::from_mut(&mut buf)); - for (dst, &src) in buf_slice.iter_mut().zip(visited[base..].iter()) { - *dst = src as u32; - } - buf - }; - let is_visited = simd.greater_than_u32s(vis_u32, zero_u); + // visited is already u32; cast directly to a SIMD vector and mask + let is_visited = simd.greater_than_u32s(*vis_v, zero_u); // r = c[j] - v[j] + (min_val - u_i) let r = simd.add_f32s(simd.sub_f32s(*c_v, *v_v), offset); @@ -207,16 +209,8 @@ impl WithSimd for InnerScan<'_> { // For argmin: mask visited lanes out with infinity let cost_for_min = simd.select_f32s_m32s(is_visited, inf_v, *spc_v); - // Column indices for this chunk - let col_indices: S::u32s = { - let mut buf = zero_u; - let buf_slice: &mut [u32] = - bytemuck::cast_slice_mut(std::slice::from_mut(&mut buf)); - for (k, dst) in buf_slice.iter_mut().enumerate() { - *dst = (base + k) as u32; - } - buf - }; + // Column indices for this chunk: iota + base + let col_indices = simd.add_u32s(iota, simd.splat_u32s(base as u32)); // Update running per-lane argmin let new_is_better = simd.less_than_f32s(cost_for_min, best_cost_v); @@ -241,7 +235,7 @@ impl WithSimd for InnerScan<'_> { // Scalar tail: remainder columns that don't fill a full SIMD vector let tail_start = ncol - c_tail.len(); for j in tail_start..ncol { - if visited[j] == 0 { + if visited[j] == 0u32 { let r = c_row[j] - u_i - v[j] + min_val; if r < spc[j] { spc[j] = r; @@ -295,11 +289,13 @@ pub fn lsap_simd(c: &[f32], nrow: usize, ncol: usize) -> Vec { let mut v = vec![0f32; ncol]; let mut col4row = vec![usize::MAX; nrow]; let mut row4col = vec![u32::MAX; ncol]; - let mut visited = vec![0u8; ncol]; + let mut visited = vec![0u32; ncol]; + let mut spc = vec![f32::INFINITY; ncol]; + let mut path = vec![u32::MAX; ncol]; for cur_row in 0..nrow { - let mut spc = vec![f32::INFINITY; ncol]; - let mut path = vec![u32::MAX; ncol]; + spc.fill(f32::INFINITY); + path.fill(u32::MAX); visited.fill(0); let mut i = cur_row; From 6abba95b269ee6683131bec02b96a38dc9a38607 Mon Sep 17 00:00:00 2001 From: gcanat Date: Thu, 2 Apr 2026 23:43:12 +0200 Subject: [PATCH 13/19] lsap core functions as a separate lib --- bindings/Cargo.lock | 14 +- powerboxesrs/Cargo.lock | 18 +- powerboxesrs/Cargo.toml | 3 +- powerboxesrs/src/assignments.rs | 390 +------------------------------- 4 files changed, 29 insertions(+), 396 deletions(-) diff --git a/bindings/Cargo.lock b/bindings/Cargo.lock index 1750f08..92df4a6 100644 --- a/bindings/Cargo.lock +++ b/bindings/Cargo.lock @@ -239,6 +239,17 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "perfect-matching" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13d16e34b987611024a80061809e2c639a7b934259a54111f0f601001343d61c" +dependencies = [ + "bytemuck", + "num-traits", + "pulp", +] + [[package]] name = "portable-atomic" version = "1.12.0" @@ -249,10 +260,9 @@ checksum = "f59e70c4aef1e55797c2e8fd94a4f2a973fc972cfde0e0b05f683667b0cd39dd" name = "powerboxesrs" version = "0.3.1" dependencies = [ - "bytemuck", "ndarray", "num-traits", - "pulp", + "perfect-matching", "rayon", "rstar", "wide", diff --git a/powerboxesrs/Cargo.lock b/powerboxesrs/Cargo.lock index 5102d78..bfb10ef 100644 --- a/powerboxesrs/Cargo.lock +++ b/powerboxesrs/Cargo.lock @@ -420,9 +420,9 @@ dependencies = [ [[package]] name = "num-traits" -version = "0.2.17" +version = "0.2.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", "libm", @@ -440,6 +440,17 @@ version = "11.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ab1bc2a289d34bd04a330323ac98a1b4bc82c9d9fcb1e66b63caa84da26b575" +[[package]] +name = "perfect-matching" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13d16e34b987611024a80061809e2c639a7b934259a54111f0f601001343d61c" +dependencies = [ + "bytemuck", + "num-traits", + "pulp", +] + [[package]] name = "plotters" version = "0.3.5" @@ -472,12 +483,11 @@ dependencies = [ name = "powerboxesrs" version = "0.3.1" dependencies = [ - "bytemuck", "codspeed-criterion-compat", "criterion", "ndarray", "num-traits", - "pulp", + "perfect-matching", "rayon", "rstar", "wide", diff --git a/powerboxesrs/Cargo.toml b/powerboxesrs/Cargo.toml index 5024255..f36ea85 100644 --- a/powerboxesrs/Cargo.toml +++ b/powerboxesrs/Cargo.toml @@ -20,10 +20,9 @@ default = ["ndarray"] ndarray = ["dep:ndarray"] [dependencies] -bytemuck = "1.15" ndarray = { version = ">=0.15, <=0.16", features = ["rayon"], optional = true } num-traits = "0.2.17" -pulp = "0.20" +perfect-matching = "0.1.0" rayon = "1.8.0" rstar = "0.11.0" wide = "0.7" diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index e2fb7e7..51e865a 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -1,358 +1,14 @@ use crate::iou::{iou_distance_slice, parallel_iou_distance_slice}; -use bytemuck::cast_slice; #[cfg(feature = "ndarray")] use ndarray::ArrayView2; use num_traits::{Num, ToPrimitive}; -use pulp::{Arch, Simd, WithSimd}; +use perfect_matching::sapjv::{lsap_scalar, lsap_simd}; // Cost matrix size above which we use parallel_iou_distance_slice const PARALLEL_IOU_MIN_BOXES: usize = 90_000; // If more than 98% boxes dont overlap, lsap_simd is slower than lsap const SIMD_MAX_SPARSITY: f64 = 0.98; -/// Shortest Augmenting Path algorithm for the Linear Sum Assignment Problem. -/// -/// Based on: Crouse, "On implementing 2D rectangular assignment algorithms", -/// https://ui.adsabs.harvard.edu/abs/2016ITAES..52.1679C/abstract -/// Given `J` jobs and `W` workers (`J <= W`), computes the minimum cost to assign each jobs -/// to distinct workers. -/// -/// # Arguments -/// * `c` - A slice representing a `J x W` cost matrix where `c[j][w]` is the cost to -/// assign job `j` to worker `w`. The slice is a row-major representation of the cost matrix. -/// -/// # Returns -/// A `Vec` of length `J`, where entry `j` is the worker's index assigned to this job. -/// -/// # Panics -/// Panics if `weights` is empty, rows have inconsistent lengths, or `J > W`. -/// -/// # Examples -/// -/// ``` -/// use powerboxesrs::assignments::lsap; -/// let costs = vec![8_i64, 5, 9, 4, 2, 4, 7, 3, 8]; -/// let assignments = lsap(&costs, 3, 3); -/// assert_eq!(assignments, vec![0, 2, 1]); -/// ``` -pub fn lsap(c: &[T], nrow: usize, ncol: usize) -> Vec -where - T: Copy + PartialOrd + std::ops::Sub + std::ops::Add, - T: num_traits::Bounded + num_traits::Zero, -{ - assert!(nrow <= ncol); - - let inf = T::max_value(); - - let mut u = vec![T::zero(); nrow]; // row potentials - let mut v = vec![T::zero(); ncol]; // col potentials - let mut col4row = vec![usize::MAX; nrow]; // col assigned to each row - let mut row4col = vec![usize::MAX; ncol]; // row assigned to each col - - for cur_row in 0..nrow { - // Dijkstra-like shortest path from cur_row to any unassigned col - let mut shortest_path_costs = vec![inf; ncol]; - let mut path = vec![usize::MAX; ncol]; - let mut visited = vec![false; ncol]; - - let mut i = cur_row; - let mut sink = usize::MAX; - let mut min_val = T::zero(); - - while sink == usize::MAX { - let mut idx = usize::MAX; - let mut lowest = inf; - - for j in 0..ncol { - if !visited[j] { - let r = c[i * ncol + j] - u[i] - v[j] + min_val; - if r < shortest_path_costs[j] { - shortest_path_costs[j] = r; - path[j] = i; - } - if shortest_path_costs[j] < lowest - || (shortest_path_costs[j] == lowest && row4col[j] == usize::MAX) - { - lowest = shortest_path_costs[j]; - idx = j; - } - } - } - - min_val = lowest; - visited[idx] = true; - - if row4col[idx] == usize::MAX { - sink = idx; - } else { - i = row4col[idx]; - } - } - - // Update potentials along the path - u[cur_row] = u[cur_row] + min_val; - for j in 0..ncol { - if visited[j] { - let r = row4col[j]; - if r != usize::MAX { - u[r] = u[r] - shortest_path_costs[j] + min_val; - } - v[j] = v[j] + shortest_path_costs[j] - min_val; - } - } - - // Augment along the path back to cur_row - let mut j = sink; - loop { - let i = path[j]; - row4col[j] = i; - let prev_j = col4row[i]; - col4row[i] = j; - if i == cur_row { - break; - } - j = prev_j; - } - } - - col4row -} -/// SIMD implementation for LSAP -struct InnerScan<'a> { - c_row: &'a [f32], - v: &'a [f32], - visited: &'a [u32], - spc: &'a mut [f32], - path: &'a mut [u32], - row4col: &'a [u32], - u_i: f32, - min_val: f32, - row_i: u32, -} - -struct ScanResult { - pub best_cost: f32, - pub best_col: usize, -} - -impl WithSimd for InnerScan<'_> { - type Output = ScanResult; - - #[inline(always)] - fn with_simd(self, simd: S) -> ScanResult { - let Self { - c_row, - v, - visited, - spc, - path, - row4col, - u_i, - min_val, - row_i, - } = self; - let ncol = c_row.len(); - - let offset = simd.splat_f32s(min_val - u_i); - let inf_v = simd.splat_f32s(f32::INFINITY); - let row_v = simd.splat_u32s(row_i); - let zero_u = simd.splat_u32s(0u32); - - let mut best_cost_v = simd.splat_f32s(f32::INFINITY); - let mut best_col_v = simd.splat_u32s(u32::MAX); - - let (c_chunks, c_tail) = S::as_simd_f32s(c_row); - let (v_chunks, _) = S::as_simd_f32s(v); - let (spc_chunks, _) = S::as_mut_simd_f32s(spc); - let (path_chunks, _) = S::as_mut_simd_u32s(path); - let (vis_chunks, _) = S::as_simd_u32s(visited); - - let lanes = std::mem::size_of::() / 4; - - // Precompute iota = [0, 1, 2, ..., lanes-1] once. - // Per-chunk indices are derived as iota + splat(base) - let iota: S::u32s = { - let mut buf = zero_u; - let buf_slice: &mut [u32] = bytemuck::cast_slice_mut(std::slice::from_mut(&mut buf)); - for (k, dst) in buf_slice.iter_mut().enumerate() { - *dst = k as u32; - } - buf - }; - - for (chunk_idx, ((((c_v, v_v), spc_v), path_v), vis_v)) in c_chunks - .iter() - .zip(v_chunks.iter()) - .zip(spc_chunks.iter_mut()) - .zip(path_chunks.iter_mut()) - .zip(vis_chunks.iter()) - .enumerate() - { - let base = chunk_idx * lanes; - - // visited is already u32; cast directly to a SIMD vector and mask - let is_visited = simd.greater_than_u32s(*vis_v, zero_u); - - // r = c[j] - v[j] + (min_val - u_i) - let r = simd.add_f32s(simd.sub_f32s(*c_v, *v_v), offset); - - // spc[j] = min(spc[j], r) for unvisited lanes only - let old_spc_v = *spc_v; - let new_spc = simd.min_f32s(*spc_v, r); - *spc_v = simd.select_f32s_m32s(is_visited, *spc_v, new_spc); - - // path[j] = row_i where r improved spc AND lane is unvisited - let improved = simd.less_than_f32s(r, old_spc_v); - let update_path = simd.and_m32s(simd.not_m32s(is_visited), improved); - *path_v = simd.select_u32s_m32s(update_path, row_v, *path_v); - - // For argmin: mask visited lanes out with infinity - let cost_for_min = simd.select_f32s_m32s(is_visited, inf_v, *spc_v); - - // Column indices for this chunk: iota + base - let col_indices = simd.add_u32s(iota, simd.splat_u32s(base as u32)); - - // Update running per-lane argmin - let new_is_better = simd.less_than_f32s(cost_for_min, best_cost_v); - best_cost_v = simd.select_f32s_m32s(new_is_better, cost_for_min, best_cost_v); - best_col_v = simd.select_u32s_m32s(new_is_better, col_indices, best_col_v); - } - - // Horizontal reduction: fold SIMD lanes down to a scalar argmin - let costs: &[f32] = cast_slice(std::slice::from_ref(&best_cost_v)); - let cols: &[u32] = cast_slice(std::slice::from_ref(&best_col_v)); - let mut best_cost = f32::INFINITY; - let mut best_col = usize::MAX; - for (&cost, &col) in costs.iter().zip(cols.iter()) { - if cost < best_cost - || (cost == best_cost && col != u32::MAX && row4col[col as usize] == u32::MAX) - { - best_cost = cost; - best_col = col as usize; - } - } - - // Scalar tail: remainder columns that don't fill a full SIMD vector - let tail_start = ncol - c_tail.len(); - for j in tail_start..ncol { - if visited[j] == 0u32 { - let r = c_row[j] - u_i - v[j] + min_val; - if r < spc[j] { - spc[j] = r; - path[j] = row_i; - } - if spc[j] < best_cost || (spc[j] == best_cost && row4col[j] == u32::MAX) { - best_cost = spc[j]; - best_col = j; - } - } - } - - ScanResult { - best_cost, - best_col, - } - } -} - -/// Shortest Augmenting Path algorithm with SIMD instructions for the Linear Sum Assignment Problem. -/// -/// Based on: Crouse, "On implementing 2D rectangular assignment algorithms", -/// https://ui.adsabs.harvard.edu/abs/2016ITAES..52.1679C/abstract -/// Given `J` jobs and `W` workers (`J <= W`), computes the minimum cost to assign each jobs -/// to distinct workers. -/// -/// # Arguments -/// * `c` - A slice representing a `J x W` cost matrix where `c[j][w]` is the cost to -/// assign job `j` to worker `w`. The slice is a row-major representation of the cost matrix. -/// -/// # Returns -/// A `Vec` of length `J`, where entry `j` is the worker's index assigned to this job. -/// -/// # Panics -/// Panics if `weights` is empty, rows have inconsistent lengths, or `J > W`. -/// -/// # Examples -/// -/// ``` -/// use powerboxesrs::assignments::lsap_simd; -/// let costs = vec![8_f32, 5., 9., 4., 2., 4., 7., 3., 8.]; -/// let assignments = lsap_simd(&costs, 3, 3); -/// assert_eq!(assignments, vec![0, 2, 1]); -/// ``` -pub fn lsap_simd(c: &[f32], nrow: usize, ncol: usize) -> Vec { - assert!(nrow <= ncol); - - let arch = Arch::new(); - - let mut u = vec![0f32; nrow]; - let mut v = vec![0f32; ncol]; - let mut col4row = vec![usize::MAX; nrow]; - let mut row4col = vec![u32::MAX; ncol]; - let mut visited = vec![0u32; ncol]; - let mut spc = vec![f32::INFINITY; ncol]; - let mut path = vec![u32::MAX; ncol]; - - for cur_row in 0..nrow { - spc.fill(f32::INFINITY); - path.fill(u32::MAX); - visited.fill(0); - - let mut i = cur_row; - let mut sink = usize::MAX; - let mut min_val = 0f32; - - while sink == usize::MAX { - let res = arch.dispatch(InnerScan { - c_row: &c[i * ncol..(i + 1) * ncol], - v: &v, - visited: &visited, - spc: &mut spc, - path: &mut path, - row4col: &row4col, - u_i: u[i], - min_val, - row_i: i as u32, - }); - - min_val = res.best_cost; - let j = res.best_col; - visited[j] = 1; - - if row4col[j] == u32::MAX { - sink = j; - } else { - i = row4col[j] as usize; - } - } - - u[cur_row] += min_val; - for j in 0..ncol { - if visited[j] != 0 { - let r = row4col[j]; - if r != u32::MAX { - u[r as usize] += min_val - spc[j]; - } - v[j] += spc[j] - min_val; - } - } - - let mut j = sink; - loop { - let pi = path[j] as usize; - row4col[j] = pi as u32; - let prev_j = col4row[pi]; - col4row[pi] = j; - if pi == cur_row { - break; - } - j = prev_j; - } - } - - col4row -} - /// Compute the optimal assignment between two sets of axis-aligned bounding boxes /// using the LSAP algorithm, minimising the total IoU distance. /// @@ -418,7 +74,7 @@ where // if cost matrix is too sparse, the overhead of calling InnerScan in the SIMD // implementation dominates so it's much faster to use the scalar implementation. let lsap_func = if non_overlapping_ratio > SIMD_MAX_SPARSITY { - lsap + lsap_scalar } else { lsap_simd }; @@ -473,48 +129,6 @@ where mod tests { use super::*; - // tests for lsap and lsap_simd - #[test] - fn test_identical_costs() { - let cost_matrix = vec![1.0_f32, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]; - let asgmt = lsap(&cost_matrix, 3, 3); - let asgmt_simd = lsap_simd(&cost_matrix, 3, 3); - // all costs are equal: matches are made starting from the end - assert_eq!(asgmt, vec![2, 1, 0]); - assert_eq!(asgmt_simd, vec![2, 1, 0]); - } - - #[test] - fn test_known_assignment() { - let costs_f32 = vec![8_f32, 5., 9., 4., 2., 4., 7., 3., 8.]; - assert_eq!(lsap(&costs_f32, 3, 3), vec![0, 2, 1]); - assert_eq!(lsap_simd(&costs_f32, 3, 3), vec![0, 2, 1]); - - let costs_i64 = vec![8_i64, 5, 9, 4, 2, 4, 7, 3, 8]; - assert_eq!(lsap(&costs_i64, 3, 3), vec![0, 2, 1]); - } - - #[test] - fn test_single_element() { - let costs = vec![42.0_f32]; - assert_eq!(lsap(&costs, 1, 1), vec![0]); - assert_eq!(lsap_simd(&costs, 1, 1), vec![0]); - } - - #[test] - fn test_rectangular() { - let costs = vec![3.0_f32, 1., 2., 1., 3., 2.]; - assert_eq!(lsap(&costs, 2, 3), vec![1, 0]); - assert_eq!(lsap_simd(&costs, 2, 3), vec![1, 0]); - } - - #[test] - fn test_optimal_in_diagonal() { - let costs = vec![0.0_f32, 1., 1., 1., 0., 1., 1., 1., 0.]; - assert_eq!(lsap(&costs, 3, 3), vec![0, 1, 2]); - assert_eq!(lsap_simd(&costs, 3, 3), vec![0, 1, 2]); - } - // tests for lsap_iou_slice #[test] fn test_identical_boxes() { From 297ff803ec2c50c28cf9501e997747436ca2c643 Mon Sep 17 00:00:00 2001 From: guillaume Date: Mon, 25 May 2026 12:53:00 +0200 Subject: [PATCH 14/19] review fixes: docstring, error propagation, python tests --- bindings/src/lib.rs | 4 +-- bindings/tests/test_dtypes.py | 47 +++++++++++++++++++++++++++++++++ powerboxesrs/src/assignments.rs | 15 +++++------ 3 files changed, 56 insertions(+), 10 deletions(-) diff --git a/bindings/src/lib.rs b/bindings/src/lib.rs index adf7e47..2117364 100644 --- a/bindings/src/lib.rs +++ b/bindings/src/lib.rs @@ -826,8 +826,8 @@ fn lsap_iou_generic( where T: Num + ToPrimitive + PartialOrd + numpy::Element + Copy, { - let boxes1 = preprocess_boxes(boxes1).unwrap(); - let boxes2 = preprocess_boxes(boxes2).unwrap(); + let boxes1 = preprocess_boxes(boxes1)?; + let boxes2 = preprocess_boxes(boxes2)?; let asgmt = assignments::lsap_iou(boxes1, boxes2, iou_threshold); Ok(asgmt) } diff --git a/bindings/tests/test_dtypes.py b/bindings/tests/test_dtypes.py index b1ee545..8eaaa1c 100644 --- a/bindings/tests/test_dtypes.py +++ b/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,49 @@ 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): + # use non-degenerate integer-friendly coordinates so the cost matrix + # is well-defined across every supported dtype + topleft = np.random.uniform(0, 500, size=(50, 2)) + wh = np.random.uniform(10, 100, size=(50, 2)) + boxes1 = np.concatenate([topleft, topleft + wh], axis=1) + topleft2 = np.random.uniform(0, 500, 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) diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index 51e865a..c2b8e38 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -12,10 +12,10 @@ const SIMD_MAX_SPARSITY: f64 = 0.98; /// Compute the optimal assignment between two sets of axis-aligned bounding boxes /// using the LSAP algorithm, minimising the total IoU distance. /// -/// Given `n1` ground-truth boxes and `n2` predicted boxes the function builds the -/// `min(n1,n2) × max(n1,n2)` cost matrix from `iou_distance_slice`, scales the -/// `f64` costs to `i64` (multiplied by `1e15`), and calls `lsap`. -/// After matching, pairs whose IoU is below `iou_threshold` are discarded. +/// Builds the `min(n1,n2) × max(n1,n2)` cost matrix from `iou_distance_slice`, +/// casts the `f64` costs to `f32`, then dispatches to `lsap_simd` or `lsap_scalar` +/// depending on cost-matrix sparsity. After matching, pairs whose IoU is below +/// `iou_threshold` are discarded. /// /// # Arguments /// @@ -27,9 +27,8 @@ const SIMD_MAX_SPARSITY: f64 = 0.98; /// /// # Returns /// -/// A pair `(indices1, indices2)` of equal-length `Vec` such that -/// `boxes1[indices1[k]]` is matched to `boxes2[indices2[k]]`. -/// The length of both vectors is at most `min(n1, n2)`. +/// A `Vec<(usize, usize)>` of matched index pairs `(i, j)` such that +/// `boxes1[i]` is matched to `boxes2[j]`. Length is at most `min(n1, n2)`. pub fn lsap_iou_slice( boxes1: &[N], boxes2: &[N], @@ -109,7 +108,7 @@ where /// /// # Returns /// -/// A pair `(indices1, indices2)` of equal-length `Vec` of length at most `min(N, M)`. +/// A `Vec<(usize, usize)>` of matched index pairs of length at most `min(N, M)`. #[cfg(feature = "ndarray")] pub fn lsap_iou<'a, N, BA>(boxes1: BA, boxes2: BA, iou_threshold: f64) -> Vec<(usize, usize)> where From 324ee390bc0782fa87fd6694f7f14a13196bc02b Mon Sep 17 00:00:00 2001 From: guillaume Date: Mon, 25 May 2026 12:57:29 +0200 Subject: [PATCH 15/19] fix: seed lsap_iou test, bound coords for uint8 cast --- bindings/tests/test_dtypes.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/bindings/tests/test_dtypes.py b/bindings/tests/test_dtypes.py index 8eaaa1c..32ce2e9 100644 --- a/bindings/tests/test_dtypes.py +++ b/bindings/tests/test_dtypes.py @@ -477,12 +477,12 @@ def test_rotated_nms_bad_dtype(): @pytest.mark.parametrize("dtype", supported_dtypes) def test_lsap_iou(dtype): - # use non-degenerate integer-friendly coordinates so the cost matrix - # is well-defined across every supported dtype - topleft = np.random.uniform(0, 500, size=(50, 2)) - wh = np.random.uniform(10, 100, size=(50, 2)) + # seeded + coords bounded so the cast is valid for every dtype, uint8 included + rng = np.random.default_rng(0) + topleft = rng.uniform(0, 150, size=(50, 2)) + wh = rng.uniform(10, 50, size=(50, 2)) boxes1 = np.concatenate([topleft, topleft + wh], axis=1) - topleft2 = np.random.uniform(0, 500, size=(50, 2)) + topleft2 = rng.uniform(0, 150, size=(50, 2)) boxes2 = np.concatenate([topleft2, topleft2 + wh], axis=1) pairs = lsap_iou(boxes1.astype(dtype), boxes2.astype(dtype)) assert isinstance(pairs, list) From f17b0ae7dc4e9ea2b66dd92519d865424fd63ab3 Mon Sep 17 00:00:00 2001 From: guillaume Date: Mon, 25 May 2026 13:01:34 +0200 Subject: [PATCH 16/19] fix: tighter wh bound in lsap_iou test to avoid u8 area overflow --- bindings/tests/test_dtypes.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/bindings/tests/test_dtypes.py b/bindings/tests/test_dtypes.py index 32ce2e9..4c65f26 100644 --- a/bindings/tests/test_dtypes.py +++ b/bindings/tests/test_dtypes.py @@ -477,12 +477,13 @@ def test_rotated_nms_bad_dtype(): @pytest.mark.parametrize("dtype", supported_dtypes) def test_lsap_iou(dtype): - # seeded + coords bounded so the cast is valid for every dtype, uint8 included + # 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, 150, size=(50, 2)) - wh = rng.uniform(10, 50, size=(50, 2)) + 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, 150, size=(50, 2)) + 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) From 76c7f46e55b34c8ca39affec76b7ce3b80c51ca9 Mon Sep 17 00:00:00 2001 From: guillaume Date: Mon, 25 May 2026 13:07:15 +0200 Subject: [PATCH 17/19] test: cover parallel_iou_distance_slice directly --- powerboxesrs/src/iou.rs | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/powerboxesrs/src/iou.rs b/powerboxesrs/src/iou.rs index 423f2ca..3263023 100644 --- a/powerboxesrs/src/iou.rs +++ b/powerboxesrs/src/iou.rs @@ -410,6 +410,42 @@ mod tests { assert_eq!(result, vec![0.0]); } + #[test] + fn test_parallel_iou_distance_slice() { + let boxes1 = vec![0.0, 0.0, 2.0, 2.0]; + let boxes2 = vec![1.0, 1.0, 3.0, 3.0]; + let result = parallel_iou_distance_slice(&boxes1, &boxes2, 1, 1); + assert_eq!(result, vec![0.8571428571428572]); + } + + #[test] + fn test_parallel_iou_distance_slice_no_overlap() { + let boxes1 = vec![0.0, 0.0, 2.0, 2.0]; + let boxes2 = vec![3.0, 3.0, 4.0, 4.0]; + let result = parallel_iou_distance_slice(&boxes1, &boxes2, 1, 1); + assert_eq!(result, vec![1.0]); + } + + #[test] + fn test_parallel_iou_distance_slice_perfect() { + let boxes1 = vec![0.0, 0.0, 2.0, 2.0]; + let boxes2 = vec![0.0, 0.0, 2.0, 2.0]; + let result = parallel_iou_distance_slice(&boxes1, &boxes2, 1, 1); + assert_eq!(result, vec![0.0]); + } + + #[test] + fn test_parallel_iou_distance_slice_matches_serial() { + // multi-row case to exercise rayon's row partitioning + let boxes1 = vec![ + 0.0, 0.0, 2.0, 2.0, 1.0, 1.0, 3.0, 3.0, 5.0, 5.0, 6.0, 6.0, + ]; + let boxes2 = vec![0.0, 0.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0]; + let serial = iou_distance_slice(&boxes1, &boxes2, 3, 2); + let parallel = parallel_iou_distance_slice(&boxes1, &boxes2, 3, 2); + assert_eq!(serial, parallel); + } + #[cfg(feature = "ndarray")] mod ndarray_tests { use super::*; From debde194add460ae89af17400e4fbc2ac13651c1 Mon Sep 17 00:00:00 2001 From: guillaume Date: Mon, 25 May 2026 13:09:01 +0200 Subject: [PATCH 18/19] fmt --- powerboxesrs/src/iou.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/powerboxesrs/src/iou.rs b/powerboxesrs/src/iou.rs index 3263023..9ec7ab2 100644 --- a/powerboxesrs/src/iou.rs +++ b/powerboxesrs/src/iou.rs @@ -437,9 +437,7 @@ mod tests { #[test] fn test_parallel_iou_distance_slice_matches_serial() { // multi-row case to exercise rayon's row partitioning - let boxes1 = vec![ - 0.0, 0.0, 2.0, 2.0, 1.0, 1.0, 3.0, 3.0, 5.0, 5.0, 6.0, 6.0, - ]; + let boxes1 = vec![0.0, 0.0, 2.0, 2.0, 1.0, 1.0, 3.0, 3.0, 5.0, 5.0, 6.0, 6.0]; let boxes2 = vec![0.0, 0.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0]; let serial = iou_distance_slice(&boxes1, &boxes2, 3, 2); let parallel = parallel_iou_distance_slice(&boxes1, &boxes2, 3, 2); From 888edee2f7c0bc1a067317dd0dee493982dfae45 Mon Sep 17 00:00:00 2001 From: guillaume Date: Mon, 25 May 2026 13:15:00 +0200 Subject: [PATCH 19/19] test: cover parallel iou branch in lsap_iou_slice --- powerboxesrs/src/assignments.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs index c2b8e38..e5fc028 100644 --- a/powerboxesrs/src/assignments.rs +++ b/powerboxesrs/src/assignments.rs @@ -211,6 +211,24 @@ mod tests { assert_eq!(pairs, vec![(0, 1), (1, 0)]); } + #[test] + fn test_parallel_iou_branch() { + // n1 * n2 = 90_601 > PARALLEL_IOU_MIN_BOXES (90_000), so this exercises + // the parallel_iou_distance_slice arm. Identity boxes keep the assignment + // trivial: pair i with i for all i. + let n = 301; + let mut boxes = Vec::with_capacity(n * 4); + for i in 0..n { + let f = i as f64; + boxes.extend_from_slice(&[f, f, f + 1.0, f + 1.0]); + } + let pairs = lsap_iou_slice(&boxes, &boxes, n, n, 0.0); + assert_eq!(pairs.len(), n); + for pair in pairs.iter() { + assert_eq!(pair.0, pair.1); + } + } + #[test] fn test_no_overlap_no_match() { let gt = vec![0.0_f64, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0];