feat: mask to boxes (#20)

Smirkey · Dec 4, 2023 · 446ccb6 · 446ccb6
1 parent 03270eb
commit 446ccb6
Show file tree

Hide file tree

Showing 10 changed files with 181 additions and 27 deletions.
diff --git a/bindings/python/powerboxes/__init__.py b/bindings/python/powerboxes/__init__.py
@@ -13,6 +13,8 @@
     _dtype_to_func_iou_distance,
     _dtype_to_func_parallel_iou_distance,
 )
+from ._powerboxes import masks_to_boxes as _masks_to_boxes
+
 from typing import TypeVar, Union
 BOXES_NOT_SAME_TYPE = "boxes1 and boxes2 must have the same dtype"
 BOXES_NOT_NP_ARRAY = "boxes must be numpy array"
@@ -175,6 +177,23 @@ def box_convert(boxes: npt.NDArray[T], in_fmt: str, out_fmt: str) -> npt.NDArray
         raise TypeError(BOXES_NOT_NP_ARRAY)
     return _dtype_to_func_box_convert[boxes.dtype](boxes, in_fmt, out_fmt)
 
+def masks_to_boxes(masks: npt.NDArray[np.bool_]) -> npt.NDArray[np.uint64]:
+    """
+    Compute the bounding boxes around the provided masks.
+
+    Returns a [N, 4] tensor containing bounding boxes. The boxes are in ``(x1, y1, x2, y2)`` format with
+    ``0 <= x1 < x2`` and ``0 <= y1 < y2``.
+
+    Args:
+        masks (Tensor[N, H, W]): masks to transform where N is the number of masks
+            and (H, W) are the spatial dimensions.
+
+    Returns:
+        Tensor[N, 4]: bounding boxes
+    """
+    if not isinstance(masks, np.ndarray):
+        raise TypeError(BOXES_NOT_NP_ARRAY)
+    return _masks_to_boxes(masks)
 
 __all__ = [
     "iou_distance",
@@ -184,5 +203,6 @@ def box_convert(boxes: npt.NDArray[T], in_fmt: str, out_fmt: str) -> npt.NDArray
     "box_convert",
     "giou_distance",
     "parallel_giou_distance",
+    "masks_to_boxes",
     "supported_dtypes"
 ]
diff --git a/bindings/src/lib.rs b/bindings/src/lib.rs
@@ -1,10 +1,10 @@
 mod utils;
 
 use num_traits::{Num, ToPrimitive};
-use numpy::{PyArray1, PyArray2};
+use numpy::{PyArray1, PyArray2, PyArray3};
 use powerboxesrs::{boxes, giou, iou};
 use pyo3::prelude::*;
-use utils::preprocess_array;
+use utils::{preprocess_array3, preprocess_boxes};
 
 #[pymodule]
 fn _powerboxes(_py: Python, m: &PyModule) -> PyResult<()> {
@@ -78,8 +78,18 @@ fn _powerboxes(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(box_convert_u32, m)?)?;
     m.add_function(wrap_pyfunction!(box_convert_u16, m)?)?;
     m.add_function(wrap_pyfunction!(box_convert_u8, m)?)?;
+    // Masks to boxes
+    m.add_function(wrap_pyfunction!(masks_to_boxes, m)?)?;
     Ok(())
 }
+// Masks to boxes
+#[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_as_numpy = utils::array_to_numpy(_py, boxes).unwrap();
+    return Ok(boxes_as_numpy.to_owned());
+}
 
 // IoU
 fn iou_distance_generic<T>(
@@ -90,8 +100,8 @@ fn iou_distance_generic<T>(
 where
     T: Num + ToPrimitive + PartialOrd + numpy::Element + Copy,
 {
-    let boxes1 = preprocess_array(boxes1).unwrap();
-    let boxes2 = preprocess_array(boxes2).unwrap();
+    let boxes1 = preprocess_boxes(boxes1).unwrap();
+    let boxes2 = preprocess_boxes(boxes2).unwrap();
     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());
@@ -178,8 +188,8 @@ fn parallel_iou_distance_generic<T>(
 where
     T: Num + ToPrimitive + PartialOrd + numpy::Element + Copy + Sync + Send,
 {
-    let boxes1 = preprocess_array(boxes1).unwrap();
-    let boxes2 = preprocess_array(boxes2).unwrap();
+    let boxes1 = preprocess_boxes(boxes1).unwrap();
+    let boxes2 = preprocess_boxes(boxes2).unwrap();
     let iou = iou::parallel_iou_distance(&boxes1, &boxes2);
     let iou_as_numpy = utils::array_to_numpy(_py, iou).unwrap();
     return Ok(iou_as_numpy.to_owned());
@@ -265,8 +275,8 @@ fn giou_distance_generic<T>(
 where
     T: Num + ToPrimitive + PartialOrd + numpy::Element + Copy,
 {
-    let boxes1 = preprocess_array(boxes1).unwrap();
-    let boxes2 = preprocess_array(boxes2).unwrap();
+    let boxes1 = preprocess_boxes(boxes1).unwrap();
+    let boxes2 = preprocess_boxes(boxes2).unwrap();
     let iou = giou::giou_distance(&boxes1, &boxes2);
     let iou_as_numpy = utils::array_to_numpy(_py, iou).unwrap();
     return Ok(iou_as_numpy.to_owned());
@@ -352,8 +362,8 @@ fn parallel_giou_distance_generic<T>(
 where
     T: Num + ToPrimitive + PartialOrd + numpy::Element + Copy,
 {
-    let boxes1 = preprocess_array(boxes1).unwrap();
-    let boxes2 = preprocess_array(boxes2).unwrap();
+    let boxes1 = preprocess_boxes(boxes1).unwrap();
+    let boxes2 = preprocess_boxes(boxes2).unwrap();
     let iou = giou::giou_distance(&boxes1, &boxes2);
     let iou_as_numpy = utils::array_to_numpy(_py, iou).unwrap();
     return Ok(iou_as_numpy.to_owned());
@@ -439,7 +449,7 @@ fn remove_small_boxes_generic<T>(
 where
     T: Num + ToPrimitive + PartialOrd + numpy::Element + Copy,
 {
-    let boxes = preprocess_array(boxes).unwrap();
+    let boxes = preprocess_boxes(boxes).unwrap();
     let filtered_boxes = boxes::remove_small_boxes(&boxes, min_size);
     let filtered_boxes_as_numpy = utils::array_to_numpy(_py, filtered_boxes).unwrap();
     return Ok(filtered_boxes_as_numpy.to_owned());
@@ -521,7 +531,7 @@ fn generic_box_areas<T>(_py: Python, boxes: &PyArray2<T>) -> PyResult<Py<PyArray
 where
     T: Num + numpy::Element + PartialOrd + ToPrimitive + Sync + Send + Copy,
 {
-    let boxes = preprocess_array(boxes).unwrap();
+    let boxes = preprocess_boxes(boxes).unwrap();
     let areas = boxes::box_areas(&boxes);
     let areas_as_numpy = utils::array_to_numpy(_py, areas).unwrap();
     return Ok(areas_as_numpy.to_owned());
@@ -573,7 +583,7 @@ fn box_convert_generic<T>(
 where
     T: Num + numpy::Element + PartialOrd + ToPrimitive + Sync + Send + Copy,
 {
-    let boxes = preprocess_array(boxes).unwrap();
+    let boxes = preprocess_boxes(boxes).unwrap();
     let in_fmt = match in_fmt {
         "xyxy" => boxes::BoxFormat::XYXY,
         "xywh" => boxes::BoxFormat::XYWH,

diff --git a/bindings/src/utils.rs b/bindings/src/utils.rs
@@ -1,6 +1,6 @@
-use ndarray::{ArrayBase, Dim, OwnedRepr};
+use ndarray::{Array2, Array3, ArrayBase, OwnedRepr};
 use num_traits::Num;
-use numpy::{IntoPyArray, PyArray, PyArray2};
+use numpy::{IntoPyArray, PyArray, PyArray2, PyArray3};
 use pyo3::prelude::*;
 
 pub fn array_to_numpy<T: numpy::Element, D: ndarray::Dimension>(
@@ -11,14 +11,11 @@ pub fn array_to_numpy<T: numpy::Element, D: ndarray::Dimension>(
     return Ok(numpy_array);
 }
 
-pub fn preprocess_array<N>(
-    array: &PyArray2<N>,
-) -> Result<ArrayBase<OwnedRepr<N>, Dim<[usize; 2]>>, PyErr>
+pub fn preprocess_boxes<N>(array: &PyArray2<N>) -> Result<Array2<N>, PyErr>
 where
     N: Num + numpy::Element + Send,
 {
-    // Usage:
-    let array: ArrayBase<OwnedRepr<N>, Dim<[usize; 2]>> = unsafe { array.as_array().to_owned() };
+    let array = unsafe { array.as_array().to_owned() };
     let array_shape = array.shape();
 
     if array_shape[1] != 4 {
@@ -42,6 +39,14 @@ where
     return Ok(array);
 }
 
+pub fn preprocess_array3<N>(array: &PyArray3<N>) -> Array3<N>
+where
+    N: numpy::Element,
+{
+    let array = unsafe { array.as_array().to_owned() };
+    return array;
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -59,20 +64,20 @@ mod tests {
     }
 
     #[test]
-    fn test_preprocess_array() {
+    fn test_preprocess_boxes() {
         Python::with_gil(|python| {
             let array = PyArray2::<f32>::zeros(python, [2, 4], false);
-            let result = preprocess_array::<f32>(array);
+            let result = preprocess_boxes::<f32>(array);
             assert!(result.is_ok());
             let unwrapped_result = result.unwrap();
             assert_eq!(unwrapped_result.shape(), &[2, 4]);
         });
     }
     #[test]
-    fn test_preprocess_array_bad_shape() {
+    fn test_preprocess_boxes_bad_shape() {
         Python::with_gil(|python| {
             let array = PyArray2::<f32>::zeros(python, [2, 16], false);
-            let result = preprocess_array::<f32>(array);
+            let result = preprocess_boxes::<f32>(array);
             assert!(result.is_err());
         });
     }

diff --git a/bindings/tests/assets/masks.tiff b/bindings/tests/assets/masks.tiff
diff --git a/bindings/tests/conftest.py b/bindings/tests/conftest.py
@@ -0,0 +1,9 @@
+import pytest
+import numpy as np
+
+@pytest.fixture
+def generate_boxes(n_boxes=100):
+    im_size = 10_000
+    topleft = np.random.uniform(0.0, high=im_size, size=(n_boxes, 2))
+    wh = np.random.uniform(15, 45, size=topleft.shape)
+    return np.concatenate([topleft, topleft + wh], axis=1).astype(np.float64)
diff --git a/bindings/tests/requirements.txt b/bindings/tests/requirements.txt
@@ -2,3 +2,4 @@ pytest==7.4.3
 pytest-codspeed>=2.0.0
 numpy
 pytest-cov
+pillow
diff --git a/bindings/tests/test_boxes.py b/bindings/tests/test_boxes.py
@@ -0,0 +1,27 @@
+from powerboxes import masks_to_boxes
+import numpy as np
+import os
+from PIL import Image
+
+def test_masks_box():
+    expected = np.array(
+        [
+            [127, 2, 165, 40],
+            [2, 50, 44, 92],
+            [56, 63, 98, 100],
+            [139, 68, 175, 104],
+            [160, 112, 198, 145],
+            [49, 138, 99, 182],
+            [108, 148, 152, 213],
+        ],
+    )
+    assets_directory = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
+    mask_path = os.path.join(assets_directory, "masks.tiff")
+    image = Image.open(mask_path)
+    masks = np.zeros((image.n_frames, image.height, image.width))
+    for index in range(image.n_frames):
+        image.seek(index)
+        masks[index] = np.array(image)
+    out = masks_to_boxes(masks.astype(np.bool_))
+    assert out.dtype == np.dtype("uint64")
+    np.testing.assert_allclose(out, expected, atol=1e-4)
diff --git a/bindings/tests/test_dtypes.py b/bindings/tests/test_dtypes.py
@@ -10,7 +10,8 @@
     parallel_giou_distance,
     parallel_iou_distance,
     remove_small_boxes,
-    supported_dtypes
+    masks_to_boxes,
+    supported_dtypes,
 )
 
 np.random.seed(42)
@@ -123,3 +124,7 @@ def test_box_convert(dtype):
 def test_box_convert_bad_inputs():
     with pytest.raises(TypeError, match=BOXES_NOT_NP_ARRAY):
         box_convert("foo", "xyxy", "xywh")
+
+def test_masks_to_boxes_bad_inputs():
+    with pytest.raises(TypeError, match=BOXES_NOT_NP_ARRAY):
+        masks_to_boxes("foo")
diff --git a/bindings/tests/test_speed.py b/bindings/tests/test_speed.py
@@ -8,6 +8,7 @@
     parallel_giou_distance,
     parallel_iou_distance,
     remove_small_boxes,
+    masks_to_boxes,
     supported_dtypes
 )
 
@@ -99,3 +100,8 @@ def test_box_convert_xywh_cxcywh(benchmark, dtype):
 def test_box_convert_xywh_xyxy(benchmark, dtype):
     boxes = np.random.random((100, 4)).astype(dtype)
     benchmark(box_convert, boxes, "xywh", "xyxy")
+
+@pytest.mark.benchmark(group="masks_to_boxes")
+def test_masks_to_boxes(benchmark):
+    masks = np.array([True]*(100*100*100)).reshape((100, 100, 100))
+    benchmark(masks_to_boxes, masks)
diff --git a/powerboxesrs/src/boxes.rs b/powerboxesrs/src/boxes.rs
@@ -1,4 +1,4 @@
-use ndarray::{Array1, Array2, Axis, Zip};
+use ndarray::{Array1, Array2, Array3, Axis, Zip};
 use num_traits::{Num, ToPrimitive};
 pub enum BoxFormat {
     XYXY,
@@ -344,10 +344,70 @@ where
     return converted_boxes;
 }
 
+/// Compute the bounding boxes around the provided masks.
+/// Returns a [N, 4] array containing bounding boxes. The boxes are in xyxy format
+///
+/// # Arguments
+///
+/// * `masks` - A [N, H, W] array of masks to transform where N is the number of masks and (H, W) are the spatial dimensions of the image.
+///
+/// # Returns
+///
+/// A [N, 4] array of boxes in xyxy format.
+/// # Example
+/// ```
+/// use ndarray::{arr3, array};
+/// use powerboxesrs::boxes::masks_to_boxes;
+/// let masks = arr3(&[
+///   [[true, true, true], [false, false, false]],
+///   [[false, false, false], [true, true, true]],
+///   [[false, false, false], [false, false, true]],
+/// ]);
+/// let boxes = masks_to_boxes(&masks);
+/// assert_eq!(boxes, array![[0, 0, 2, 0], [0, 1, 2, 1], [2, 1, 2, 1]]);
+pub fn masks_to_boxes(masks: &Array3<bool>) -> Array2<usize> {
+    let num_masks = masks.shape()[0];
+    let height = masks.shape()[1];
+    let width = masks.shape()[2];
+    let mut boxes = Array2::<usize>::zeros((num_masks, 4));
+
+    for (i, mask) in masks.outer_iter().enumerate() {
+        let mut x1 = width;
+        let mut y1 = height;
+        let mut x2 = 0;
+        let mut y2 = 0;
+
+        // get the indices where the mask is true
+        mask.indexed_iter().for_each(|(index, &value)| {
+            if value {
+                let (y, x) = index;
+                if x < x1 {
+                    x1 = x;
+                }
+                if x > x2 {
+                    x2 = x;
+                }
+                if y < y1 {
+                    y1 = y;
+                }
+                if y > y2 {
+                    y2 = y;
+                }
+            }
+        });
+        boxes[[i, 0]] = x1;
+        boxes[[i, 1]] = y1;
+        boxes[[i, 2]] = x2;
+        boxes[[i, 3]] = y2;
+    }
+
+    return boxes;
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
-    use ndarray::{arr2, array};
+    use ndarray::{arr2, arr3, array};
     #[test]
     fn test_box_convert_xyxy_to_xywh() {
         let boxes = arr2(&[
@@ -529,4 +589,15 @@ mod tests {
         let filtered_boxes = remove_small_boxes(&boxes, min_size);
         assert_eq!(filtered_boxes, array![[0., 0., 10., 10.]]);
     }
+
+    #[test]
+    fn test_masks_to_boxes() {
+        let masks: Array3<bool> = arr3(&[
+            [[true, true, true], [false, false, false]],
+            [[false, false, false], [true, true, true]],
+            [[false, false, false], [false, false, true]],
+        ]);
+        let boxes = masks_to_boxes(&masks);
+        assert_eq!(boxes, array![[0, 0, 2, 0], [0, 1, 2, 1], [2, 1, 2, 1]]);
+    }
 }