diff --git a/bindings/Cargo.lock b/bindings/Cargo.lock index f04f1d3..92df4a6 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", ] @@ -202,9 +203,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", @@ -238,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" @@ -246,10 +258,11 @@ checksum = "f59e70c4aef1e55797c2e8fd94a4f2a973fc972cfde0e0b05f683667b0cd39dd" [[package]] name = "powerboxesrs" -version = "0.3.0" +version = "0.3.1" dependencies = [ "ndarray", "num-traits", + "perfect-matching", "rayon", "rstar", "wide", @@ -264,6 +277,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 +385,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" @@ -403,9 +434,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..51a0b86 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_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__", ] diff --git a/bindings/python/powerboxes/_dispatch.py b/bindings/python/powerboxes/_dispatch.py index 5bb161c..0e1d490 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_lsap_iou = _build_dispatch("lsap_iou") diff --git a/bindings/src/lib.rs b/bindings/src/lib.rs index 043fed4..2117364 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::{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> { + $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,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( + boxes1: &Bound<'_, PyArray2>, + boxes2: &Bound<'_, PyArray2>, + iou_threshold: f64, +) -> PyResult> +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); diff --git a/bindings/tests/test_dtypes.py b/bindings/tests/test_dtypes.py index b1ee545..4c65f26 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,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) diff --git a/powerboxesrs/Cargo.lock b/powerboxesrs/Cargo.lock index c59a141..bfb10ef 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", ] @@ -419,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", @@ -439,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" @@ -475,6 +487,7 @@ dependencies = [ "criterion", "ndarray", "num-traits", + "perfect-matching", "rayon", "rstar", "wide", @@ -489,6 +502,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 +549,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..f36ea85 100644 --- a/powerboxesrs/Cargo.toml +++ b/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" diff --git a/powerboxesrs/src/assignments.rs b/powerboxesrs/src/assignments.rs new file mode 100644 index 0000000..e5fc028 --- /dev/null +++ b/powerboxesrs/src/assignments.rs @@ -0,0 +1,259 @@ +use crate::iou::{iou_distance_slice, parallel_iou_distance_slice}; +#[cfg(feature = "ndarray")] +use ndarray::ArrayView2; +use num_traits::{Num, ToPrimitive}; +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; + +/// Compute the optimal assignment between two sets of axis-aligned bounding boxes +/// using the LSAP algorithm, minimising the total IoU distance. +/// +/// 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 +/// +/// * `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 `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], + 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 > PARALLEL_IOU_MIN_BOXES { + parallel_iou_distance_slice + } else { + iou_distance_slice + }; + + // 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, nrows, ncols) + } else { + iou_func(boxes2, boxes1, 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() + .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. + let lsap_func = if non_overlapping_ratio > SIMD_MAX_SPARSITY { + lsap_scalar + } else { + lsap_simd + }; + + let assignments = lsap_func(&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 LSAP algorithm, minimising the total IoU distance. +/// +/// Wraps [`lsap_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 `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 + 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"); + lsap_iou_slice(s1, s2, n1, n2, iou_threshold) +} + +#[cfg(test)] +mod tests { + use super::*; + + // 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]; + 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); + } + } + + #[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 = lsap_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 = lsap_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 = lsap_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 = lsap_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 = lsap_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 = lsap_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_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]; + 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 = lsap_iou_slice(>, &pred, 2, 2, 0.1); + assert!( + pairs.is_empty(), + "expected no matches when IoU is 0 for all pairs, got {:?}", + pairs + ); + } + + // test for ndarray lsap_iou + #[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 = lsap_iou(&boxes1, &boxes2, 0.0); + assert_eq!(pairs, vec![(0, 0), (1, 1)]); + } + } +} diff --git a/powerboxesrs/src/iou.rs b/powerboxesrs/src/iou.rs index de0ade2..9ec7ab2 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. @@ -397,6 +410,40 @@ 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::*; 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;