merge

Smirkey · May 27, 2024 · a5667bc · a5667bc
2 parents a5acd47 + f9e08fd
commit a5667bc
Show file tree

Hide file tree

Showing 13 changed files with 413 additions and 288 deletions.
diff --git a/.github/workflows/codespeed.yml b/.github/workflows/codespeed.yml
@@ -40,7 +40,7 @@ jobs:
         run: cd powerboxesrs && cargo codspeed build
 
       - name: Run benchmarks
-        uses: CodSpeedHQ/action@v1
+        uses: CodSpeedHQ/action@v2
         with:
           token: ${{ secrets.CODSPEED_TOKEN }}
           run: pytest bindings/tests/ --codspeed && cd powerboxesrs && cargo codspeed run
diff --git a/bindings/python/powerboxes/__init__.py b/bindings/python/powerboxes/__init__.py
@@ -136,7 +136,7 @@ def parallel_iou_distance(
         raise TypeError(_BOXES_NOT_NP_ARRAY)
     if boxes1.dtype == boxes2.dtype:
         try:
-            return _dtype_to_func_parallel_iou_distance[boxes1.dtype](boxes1, boxes2)
+            return _dtype_to_func_parallel_iou_distance[boxes1.dtype](boxes1, boxes2) # type: ignore
         except KeyError:
             raise TypeError(
                 f"Box dtype: {boxes1.dtype} not in supported dtypes {supported_dtypes}"
@@ -167,7 +167,7 @@ def parallel_giou_distance(
         raise TypeError(_BOXES_NOT_NP_ARRAY)
     if boxes1.dtype == boxes2.dtype:
         try:
-            return _dtype_to_func_parallel_giou_distance[boxes1.dtype](boxes1, boxes2)
+            return _dtype_to_func_parallel_giou_distance[boxes1.dtype](boxes1, boxes2) # type: ignore
         except KeyError:
             raise TypeError(
                 f"Box dtype: {boxes1.dtype} not in supported dtypes {supported_dtypes}"
@@ -198,7 +198,7 @@ def giou_distance(
         raise TypeError(_BOXES_NOT_NP_ARRAY)
     if boxes1.dtype == boxes2.dtype:
         try:
-            return _dtype_to_func_giou_distance[boxes1.dtype](boxes1, boxes2)
+            return _dtype_to_func_giou_distance[boxes1.dtype](boxes1, boxes2) # type: ignore
         except KeyError:
             raise TypeError(
                 f"Box dtype: {boxes1.dtype} not in supported dtypes {supported_dtypes}"
@@ -229,7 +229,7 @@ def tiou_distance(
         raise TypeError(_BOXES_NOT_NP_ARRAY)
     if boxes1.dtype == boxes2.dtype:
         try:
-            return _dtype_to_func_tiou_distance[boxes1.dtype](boxes1, boxes2)
+            return _dtype_to_func_tiou_distance[boxes1.dtype](boxes1, boxes2) # type: ignore
         except KeyError:
             raise TypeError(
                 f"Box dtype: {boxes1.dtype} not in supported dtypes {supported_dtypes}"
@@ -346,7 +346,7 @@ def remove_small_boxes(boxes: npt.NDArray[T], min_size: float) -> npt.NDArray[T]
     if not isinstance(boxes, np.ndarray):
         raise TypeError(_BOXES_NOT_NP_ARRAY)
     try:
-        return _dtype_to_func_remove_small_boxes[boxes.dtype](boxes, min_size)
+        return _dtype_to_func_remove_small_boxes[boxes.dtype](boxes, min_size) # type: ignore
     except KeyError:
         raise TypeError(
             f"Box dtype: {boxes.dtype} not in supported dtypes {supported_dtypes}"
@@ -365,7 +365,7 @@ def boxes_areas(boxes: npt.NDArray[T]) -> npt.NDArray[np.float64]:
     if not isinstance(boxes, np.ndarray):
         raise TypeError(_BOXES_NOT_NP_ARRAY)
     try:
-        return _dtype_to_func_box_areas[boxes.dtype](boxes)
+        return _dtype_to_func_box_areas[boxes.dtype](boxes) # type: ignore
     except KeyError:
         raise TypeError(
             f"Box dtype: {boxes.dtype} not in supported dtypes {supported_dtypes}"
@@ -391,7 +391,7 @@ def box_convert(boxes: npt.NDArray[T], in_fmt: str, out_fmt: str) -> npt.NDArray
     if not isinstance(boxes, np.ndarray):
         raise TypeError(_BOXES_NOT_NP_ARRAY)
     try:
-        return _dtype_to_func_box_convert[boxes.dtype](boxes, in_fmt, out_fmt)
+        return _dtype_to_func_box_convert[boxes.dtype](boxes, in_fmt, out_fmt) # type: ignore
     except KeyError:
         raise TypeError(
             f"Box dtype: {boxes.dtype} not in supported dtypes {supported_dtypes}"
@@ -438,7 +438,7 @@ def nms(
     if not isinstance(boxes, np.ndarray) or not isinstance(scores, np.ndarray):
         raise TypeError("Boxes and scores must be numpy arrays")
     try:
-        return _dtype_to_func_nms[boxes.dtype](
+        return _dtype_to_func_nms[boxes.dtype]( # type: ignore
             boxes, scores, iou_threshold, score_threshold
         )
     except KeyError:
@@ -473,7 +473,7 @@ def rtree_nms(
     if not isinstance(boxes, np.ndarray) or not isinstance(scores, np.ndarray):
         raise TypeError("Boxes and scores must be numpy arrays")
     try:
-        return _dtype_to_func_rtree_nms[boxes.dtype](
+        return _dtype_to_func_rtree_nms[boxes.dtype]( # type: ignore
             boxes, scores, iou_threshold, score_threshold
         )
     except KeyError:

diff --git a/bindings/src/lib.rs b/bindings/src/lib.rs
@@ -2,6 +2,7 @@ mod utils;
 
 use std::fmt::Debug;
 
+use ndarray::Array1;
 use num_traits::{Bounded, Float, Num, Signed, ToPrimitive};
 use numpy::{PyArray1, PyArray2, PyArray3};
 use powerboxesrs::{boxes, diou, giou, iou, nms, tiou};
@@ -123,7 +124,7 @@ fn _powerboxes(_py: Python, m: &PyModule) -> PyResult<()> {
 #[pyfunction]
 fn masks_to_boxes(_py: Python, masks: &PyArray3<bool>) -> PyResult<Py<PyArray2<usize>>> {
     let masks = preprocess_array3(masks);
-    let boxes = boxes::masks_to_boxes(&masks);
+    let boxes = boxes::masks_to_boxes(masks);
     let boxes_as_numpy = utils::array_to_numpy(_py, boxes).unwrap();
     return Ok(boxes_as_numpy.to_owned());
 }
@@ -179,7 +180,7 @@ fn diou_distance_generic<T>(
     boxes2: &PyArray2<T>,
 ) -> PyResult<Py<PyArray2<f64>>>
 where
-    T: Float + numpy::Element,
+    T: Num + Float + numpy::Element,
 {
     let boxes1 = preprocess_boxes(boxes1).unwrap();
     let boxes2 = preprocess_boxes(boxes2).unwrap();
@@ -216,7 +217,7 @@ where
 {
     let boxes1 = preprocess_boxes(boxes1).unwrap();
     let boxes2 = preprocess_boxes(boxes2).unwrap();
-    let iou = iou::iou_distance(&boxes1, &boxes2);
+    let iou = iou::iou_distance(boxes1, boxes2);
     let iou_as_numpy = utils::array_to_numpy(_py, iou).unwrap();
     return Ok(iou_as_numpy.to_owned());
 }
@@ -806,7 +807,7 @@ where
             ))
         }
     };
-    let converted_boxes = boxes::box_convert(&boxes, &in_fmt, &out_fmt);
+    let converted_boxes = boxes::box_convert(&boxes, in_fmt, out_fmt);
     let converted_boxes_as_numpy = utils::array_to_numpy(_py, converted_boxes).unwrap();
     return Ok(converted_boxes_as_numpy.to_owned());
 }
@@ -902,12 +903,13 @@ fn nms_generic<T>(
     score_threshold: f64,
 ) -> PyResult<Py<PyArray1<usize>>>
 where
-    T: Num + numpy::Element + PartialOrd + ToPrimitive + Copy,
+    T: numpy::Element + Num + PartialEq + PartialOrd + ToPrimitive + Copy,
 {
     let boxes = preprocess_boxes(boxes).unwrap();
     let scores = preprocess_array1(scores);
     let keep = nms::nms(&boxes, &scores, iou_threshold, score_threshold);
-    let keep_as_numpy = utils::array_to_numpy(_py, keep).unwrap();
+    let keep_as_ndarray = Array1::from(keep);
+    let keep_as_numpy = utils::array_to_numpy(_py, keep_as_ndarray).unwrap();
     return Ok(keep_as_numpy.to_owned());
 }
 #[pyfunction]
@@ -1064,21 +1066,23 @@ fn rtree_nms_generic<T>(
     score_threshold: f64,
 ) -> PyResult<Py<PyArray1<usize>>>
 where
-    T: Num
-        + numpy::Element
-        + PartialOrd
-        + ToPrimitive
-        + Copy
+    T: numpy::Element
+        + Num
         + Signed
         + Bounded
         + Debug
+        + PartialEq
+        + PartialOrd
+        + ToPrimitive
+        + Copy
         + Sync
         + Send,
 {
     let boxes = preprocess_boxes(boxes).unwrap();
     let scores = preprocess_array1(scores);
     let keep = nms::rtree_nms(&boxes, &scores, iou_threshold, score_threshold);
-    let keep_as_numpy = utils::array_to_numpy(_py, keep).unwrap();
+    let keep_as_ndarray = Array1::from(keep);
+    let keep_as_numpy = utils::array_to_numpy(_py, keep_as_ndarray).unwrap();
     return Ok(keep_as_numpy.to_owned());
 }
 #[pyfunction]

diff --git a/bindings/src/utils.rs b/bindings/src/utils.rs
@@ -1,19 +1,44 @@
-use ndarray::{Array1, Array2, Array3, ArrayBase, OwnedRepr};
+use ndarray::{ArrayBase, Dim, OwnedRepr, ViewRepr};
 use num_traits::Num;
 use numpy::{IntoPyArray, PyArray, PyArray1, PyArray2, PyArray3};
 use pyo3::prelude::*;
 
-pub fn array_to_numpy<T: numpy::Element, D: ndarray::Dimension>(
+/// Converts a 2-dimensional Rust ndarray to a NumPy array.
+///
+/// # Arguments
+///
+/// * `py` - The Python interpreter context.
+/// * `array` - The 2-dimensional Rust ndarray to convert.
+///
+/// # Returns
+///
+/// A reference to the converted NumPy array.
+///
+/// # Example
+///
+/// ```rust
+/// let py = Python::acquire_gil().python();
+/// let array_2d: Array2<f64> = Array2::ones((3, 3));
+/// let numpy_array_2d = array2_to_numpy(py, array_2d).unwrap();
+/// ```
+pub fn array_to_numpy<T, D>(
     py: Python,
     array: ArrayBase<OwnedRepr<T>, D>,
-) -> PyResult<&PyArray<T, D>> {
-    let numpy_array: &PyArray<T, D> = array.into_pyarray(py);
+) -> PyResult<&PyArray<T, D>>
+where
+    T: numpy::Element,
+    D: ndarray::Dimension,
+{
+    let numpy_array = array.into_pyarray(py);
+
     return Ok(numpy_array);
 }
 
-pub fn preprocess_boxes<N>(array: &PyArray2<N>) -> Result<Array2<N>, PyErr>
+pub fn preprocess_boxes<N>(
+    array: &PyArray2<N>,
+) -> Result<ArrayBase<ViewRepr<&N>, Dim<[usize; 2]>>, PyErr>
 where
-    N: Num + numpy::Element + Send,
+    N: numpy::Element,
 {
     let array = unsafe { array.as_array() };
     let array_shape = array.shape();
@@ -32,16 +57,14 @@ where
         }
     }
 
-    let array = array
-        .to_owned()
-        .into_shape((array_shape[0], array_shape[1]))
-        .unwrap();
     return Ok(array);
 }
 
-pub fn preprocess_rotated_boxes<N>(array: &PyArray2<N>) -> Result<Array2<N>, PyErr>
+pub fn preprocess_rotated_boxes<'a, N>(
+    array: &PyArray2<N>,
+) -> Result<ArrayBase<ViewRepr<&N>, Dim<[usize; 2]>>, PyErr>
 where
-    N: Num + numpy::Element + Send,
+    N: Num + numpy::Element + Send + 'a,
 {
     let array = unsafe { array.as_array() };
     let array_shape = array.shape();
@@ -60,42 +83,37 @@ where
         }
     }
 
-    let array = array
-        .to_owned()
-        .into_shape((array_shape[0], array_shape[1]))
-        .unwrap();
     return Ok(array);
 }
 
-pub fn preprocess_array3<N>(array: &PyArray3<N>) -> Array3<N>
+pub fn preprocess_array3<'a, N>(array: &PyArray3<N>) -> ArrayBase<ViewRepr<&N>, Dim<[usize; 3]>>
 where
-    N: numpy::Element,
+    N: numpy::Element + 'a,
 {
-    let array = unsafe { array.as_array().to_owned() };
+    let array = unsafe { array.as_array() };
     return array;
 }
 
-pub fn preprocess_array1<N>(array: &PyArray1<N>) -> Array1<N>
+pub fn preprocess_array1<'a, N>(array: &PyArray1<N>) -> ArrayBase<ViewRepr<&N>, Dim<[usize; 1]>>
 where
-    N: numpy::Element,
+    N: numpy::Element + 'a,
 {
-    let array = unsafe { array.as_array().to_owned() };
+    let array: ArrayBase<ViewRepr<&N>, ndarray::prelude::Dim<[usize; 1]>> =
+        unsafe { array.as_array() };
     return array;
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use ndarray::ArrayBase;
+    use ndarray::Array1;
 
     #[test]
     fn test_array_to_numpy() {
-        let data = vec![1., 2., 3., 4.];
-        let array = ArrayBase::from_shape_vec((1, 4), data).unwrap();
+        let array = Array1::from(vec![1., 2., 3., 4.]);
         Python::with_gil(|py| {
             let result = array_to_numpy(py, array).unwrap();
-            assert_eq!(result.readonly().shape(), &[1, 4]);
-            assert_eq!(result.readonly().shape(), &[1, 4]);
+            assert_eq!(result.readonly().shape(), &[4]);
         });
     }
 

diff --git a/bindings/tests/test_boxes.py b/bindings/tests/test_boxes.py
@@ -1,7 +1,8 @@
-from powerboxes import masks_to_boxes
-import numpy as np
 import os
+
+import numpy as np
 from PIL import Image
+from powerboxes import masks_to_boxes
 
 
 def test_masks_box():