feat: rotated giou distance (#37)

README.md

@@ -47,6 +47,11 @@ intersection = pb.iou_distance(box, box)
 - `parallel_iou_distance`: Compute the intersection over union matrix of two sets of boxes in parallel
 - `giou_distance`: Compute the generalized intersection over union matrix of two sets of boxes
 - `parallel_giou_distance`: Compute the generalized intersection over union matrix of two sets of boxes in parallel
+- `tiou_distance`: Compute the tracking intersection over union matrix of two sets of boxes
+
+#### Rotated Box Metrics
+- `rotated_iou_distance`: Compute the intersection over union matrix of two sets of rotated boxes in cxcywha format
+- `rotated_giou_distance`: Compute the generalized intersection over union matrix of two sets of rotated boxes in cxcywha format
 
 #### Box NMS
 - `nms`: Non-maximum suppression, returns the indices of the boxes to keep

bindings/README.md

@@ -40,6 +40,11 @@ intersection = pb.iou_distance(box, box)
 - `parallel_iou_distance`: Compute the intersection over union matrix of two sets of boxes in parallel
 - `giou_distance`: Compute the generalized intersection over union matrix of two sets of boxes
 - `parallel_giou_distance`: Compute the generalized intersection over union matrix of two sets of boxes in parallel
+- `tiou_distance`: Compute the tracking intersection over union matrix of two sets of boxes
+
+#### Rotated Box Metrics
+- `rotated_iou_distance`: Compute the intersection over union matrix of two sets of rotated boxes in cxcywha format
+- `rotated_giou_distance`: Compute the generalized intersection over union matrix of two sets of rotated boxes in cxcywha format
 
 #### Box NMS
 - `nms`: Non-maximum suppression, returns the indices of the boxes to keep

bindings/python/powerboxes/__init__.py

@@ -18,6 +18,7 @@
 )
 from ._nms import _dtype_to_func_nms, _dtype_to_func_rtree_nms
 from ._powerboxes import masks_to_boxes as _masks_to_boxes
+from ._powerboxes import rotated_giou_distance as _rotated_giou_distance
 from ._powerboxes import rotated_iou_distance as _rotated_iou_distance
 from ._tiou import _dtype_to_func_tiou_distance
 
@@ -34,7 +35,7 @@
     "uint32",
     "uint64",
 ]
-__version__ = "0.1.3"
+__version__ = "0.2.0"
 
 T = TypeVar(
     "T",
@@ -204,8 +205,25 @@ def tiou_distance(
 
 
 def rotated_iou_distance(
-    boxes1: npt.NDArray[T], boxes2: npt.NDArray[T]
+    boxes1: npt.NDArray[np.float64], boxes2: npt.NDArray[np.float64]
 ) -> npt.NDArray[np.float64]:
+    """Compute the pairwise iou distance between rotated boxes
+
+    Boxes should be in (cx, cy, w, h, a) format
+    where cx and cy are center coordinates, w and h
+    width and height and a, the angle in degrees
+
+    Args:
+        boxes1: 2d array of boxes in cxywha format
+        boxes2: 2d array of boxes in cxywha format
+
+    Raises:
+        TypeError: if boxes1 or boxes2 are not numpy arrays
+        ValueError: if boxes1 and boxes2 have different dtypes
+
+    Returns:
+        np.ndarray: 2d matrix of pairwise distances
+    """
     if not isinstance(boxes1, np.ndarray) or not isinstance(boxes2, np.ndarray):
         raise TypeError(_BOXES_NOT_NP_ARRAY)
     if boxes1.dtype == boxes2.dtype == np.dtype("float64"):
@@ -216,6 +234,36 @@ def rotated_iou_distance(
         )
 
 
+def rotated_giou_distance(
+    boxes1: npt.NDArray[T], boxes2: npt.NDArray[T]
+) -> npt.NDArray[np.float64]:
+    """Compute the pairwise giou distance between rotated boxes
+
+    Boxes should be in (cx, cy, w, h, a) format
+    where cx and cy are center coordinates, w and h
+    width and height and a, the angle in degrees
+
+    Args:
+        boxes1: 2d array of boxes in cxywha format
+        boxes2: 2d array of boxes in cxywha format
+
+    Raises:
+        TypeError: if boxes1 or boxes2 are not numpy arrays
+        ValueError: if boxes1 and boxes2 have different dtypes
+
+    Returns:
+        np.ndarray: 2d matrix of pairwise distances
+    """
+    if not isinstance(boxes1, np.ndarray) or not isinstance(boxes2, np.ndarray):
+        raise TypeError(_BOXES_NOT_NP_ARRAY)
+    if boxes1.dtype == boxes2.dtype == np.dtype("float64"):
+        return _rotated_giou_distance(boxes1, boxes2)
+    else:
+        raise TypeError(
+            f"Boxes dtype: {boxes1.dtype}, {boxes2.dtype} not in float64 dtype"
+        )
+
+
 def remove_small_boxes(boxes: npt.NDArray[T], min_size) -> npt.NDArray[T]:
     """Remove boxes with area less than min_area.
 
@@ -380,7 +428,8 @@ def rtree_nms(
     "supported_dtypes",
     "nms",
     "tiou_distance",
-    "rotated_iou_distance"
+    "rotated_iou_distance",
+    "rotated_giou_distance",
     "rtree_nms",
     "__version__",
 ]

bindings/src/lib.rs

@@ -110,6 +110,8 @@ fn _powerboxes(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(masks_to_boxes, m)?)?;
     // Rotated IoU
     m.add_function(wrap_pyfunction!(rotated_iou_distance, m)?)?;
+    // Rotated GIoU
+    m.add_function(wrap_pyfunction!(rotated_giou_distance, m)?)?;
     Ok(())
 }
 // Masks to boxes
@@ -136,6 +138,21 @@ fn rotated_iou_distance(
     return Ok(iou_as_numpy.to_owned());
 }
 
+// Rotated box GIoU
+
+#[pyfunction]
+fn rotated_giou_distance(
+    _py: Python,
+    boxes1: &PyArray2<f64>,
+    boxes2: &PyArray2<f64>,
+) -> PyResult<Py<PyArray2<f64>>> {
+    let boxes1 = preprocess_rotated_boxes(boxes1).unwrap();
+    let boxes2 = preprocess_rotated_boxes(boxes2).unwrap();
+    let iou = giou::rotated_giou_distance(&boxes1, &boxes2);
+    let iou_as_numpy = utils::array_to_numpy(_py, iou).unwrap();
+    return Ok(iou_as_numpy.to_owned());
+}
+
 // IoU
 fn iou_distance_generic<T>(
     _py: Python,

bindings/tests/test_dtypes.py

@@ -12,6 +12,7 @@
     parallel_giou_distance,
     parallel_iou_distance,
     remove_small_boxes,
+    rotated_giou_distance,
     rotated_iou_distance,
     rtree_nms,
     supported_dtypes,
@@ -284,7 +285,7 @@ def test_rotated_iou_distance_bad_inputs():
     with pytest.raises(Exception):
         try:
             rotated_iou_distance(np.random.random((100, 4)), np.random.random((100, 4)))
-        except:   # noqa: E722
+        except:  # noqa: E722
             raise RuntimeError()
     with pytest.raises(RuntimeError):
         try:
@@ -301,3 +302,40 @@ def test_rotated_iou_distance_dtype():
             boxes1.astype(unsuported_dtype_example),
             boxes2.astype(unsuported_dtype_example),
         )
+
+
+@pytest.mark.parametrize("dtype", ["float64"])
+def test_rotated_giou_distance(dtype):
+    boxes1 = np.random.random((100, 5))
+    boxes2 = np.random.random((100, 5))
+    rotated_giou_distance(
+        boxes1.astype(dtype),
+        boxes2.astype(dtype),
+    )
+
+
+def test_rotated_giou_distance_bad_inputs():
+    with pytest.raises(TypeError, match=_BOXES_NOT_NP_ARRAY):
+        rotated_giou_distance("foo", "bar")
+    with pytest.raises(Exception):
+        try:
+            rotated_giou_distance(
+                np.random.random((100, 4)), np.random.random((100, 4))
+            )
+        except:  # noqa: E722
+            raise RuntimeError()
+    with pytest.raises(RuntimeError):
+        try:
+            rotated_giou_distance(np.random.random((0, 4)), np.random.random((100, 4)))
+        except:  # noqa: E722
+            raise RuntimeError()
+
+
+def test_rotated_giou_distance_dtype():
+    boxes1 = np.random.random((100, 5))
+    boxes2 = np.random.random((100, 5))
+    with pytest.raises(TypeError):
+        rotated_giou_distance(
+            boxes1.astype(unsuported_dtype_example),
+            boxes2.astype(unsuported_dtype_example),
+        )

bindings/tests/test_speed.py

@@ -10,6 +10,7 @@
     parallel_giou_distance,
     parallel_iou_distance,
     remove_small_boxes,
+    rotated_giou_distance,
     rotated_iou_distance,
     rtree_nms,
     supported_dtypes,
@@ -30,14 +31,22 @@ def generate_boxes(request):
     return np.concatenate([topleft, topleft + wh], axis=1).astype(np.float64)
 
 
-@pytest.mark.benchmark(group="tiou_distance")
+@pytest.mark.benchmark(group="rotated_iou_distance")
 @pytest.mark.parametrize("dtype", ["float64"])
 def test_rotated_iou_distance(benchmark, dtype):
     boxes1 = np.random.random((100, 5)).astype(dtype)
     boxes2 = np.random.random((100, 5)).astype(dtype)
     benchmark(rotated_iou_distance, boxes1, boxes2)
 
 
+@pytest.mark.benchmark(group="rotated_giou_distance")
+@pytest.mark.parametrize("dtype", ["float64"])
+def test_rotated_giou_distance(benchmark, dtype):
+    boxes1 = np.random.random((100, 5)).astype(dtype)
+    boxes2 = np.random.random((100, 5)).astype(dtype)
+    benchmark(rotated_giou_distance, boxes1, boxes2)
+
+
 @pytest.mark.benchmark(group="tiou_distance")
 @pytest.mark.parametrize("dtype", supported_dtypes)
 def test_tiou_distance(benchmark, dtype):

powerboxesrs/README.md

@@ -23,6 +23,11 @@ cargo add powerboxesrs
 - `parallel_iou_distance`: Compute the intersection over union matrix of two sets of boxes in parallel
 - `giou_distance`: Compute the generalized intersection over union matrix of two sets of boxes
 - `parallel_giou_distance`: Compute the generalized intersection over union matrix of two sets of boxes in parallel
+- `tiou_distance`: Compute the tracking intersection over union matrix of two sets of boxes
+
+#### Rotated Box Metrics
+- `rotated_iou_distance`: Compute the intersection over union matrix of two sets of rotated boxes in cxcywha format
+- `rotated_giou_distance`: Compute the generalized intersection over union matrix of two sets of rotated boxes in cxcywha format
 
 #### Box NMS
 - `nms`: Non-maximum suppression, returns the indices of the boxes to keep

powerboxesrs/src/giou.rs

@@ -1,7 +1,12 @@
 use ndarray::{Array2, Zip};
 use num_traits::{Num, ToPrimitive};
+use rstar::RTree;
 
-use crate::{boxes, utils};
+use crate::{
+    boxes::{self, rotated_box_areas},
+    rotation::{intersection_area, minimal_bounding_rect, Rect},
+    utils,
+};
 /// Computes the Generalized Intersection over Union (GIOU) distance between two sets of bounding boxes.
 /// # Arguments
 ///
@@ -161,6 +166,97 @@ where
     giou_matrix
 }
 
+/// Calculates the rotated Generalized IoU (Giou) distance between two sets of rotated bounding boxes.
+///
+/// Given two sets of rotated bounding boxes represented by `boxes1` and `boxes2`, this function
+/// computes the rotated Giou distance matrix between them. The rotated Giou distance is a measure
+/// of dissimilarity between two rotated bounding boxes, taking into account both their overlap
+/// and the encompassing area.
+///
+/// # Arguments
+///
+/// * `boxes1` - A reference to a 2D array (Array2) containing the parameters of the first set of rotated bounding boxes.
+/// Each row of `boxes1` represents a rotated bounding box with parameters [center_x, center_y, width, height, angle].
+///
+/// * `boxes2` - A reference to a 2D array (Array2) containing the parameters of the second set of rotated bounding boxes.
+/// Each row of `boxes2` represents a rotated bounding box with parameters [center_x, center_y, width, height, angle].
+///
+/// # Returns
+///
+/// A 2D array (Array2) representing the rotated Giou distance matrix between the input sets of rotated bounding boxes.
+/// The element at position (i, j) in the matrix represents the rotated Giou distance between the i-th box in `boxes1` and
+/// the j-th box in `boxes2`.
+///
+pub fn rotated_giou_distance(boxes1: &Array2<f64>, boxes2: &Array2<f64>) -> Array2<f64> {
+    let num_boxes1 = boxes1.nrows();
+    let num_boxes2 = boxes2.nrows();
+
+    let mut iou_matrix = Array2::<f64>::ones((num_boxes1, num_boxes2));
+    let areas1 = rotated_box_areas(&boxes1);
+    let areas2 = rotated_box_areas(&boxes2);
+
+    let boxes1_rects: Vec<Rect> = boxes1
+        .rows()
+        .into_iter()
+        .map(|row| Rect::new(row[0], row[1], row[2], row[3], row[4]))
+        .collect();
+    let boxes2_rects: Vec<Rect> = boxes2
+        .rows()
+        .into_iter()
+        .map(|row| Rect::new(row[0], row[1], row[2], row[3], row[4]))
+        .collect();
+    let boxes1_bounding_rects: Vec<utils::Bbox<f64>> = boxes1_rects
+        .iter()
+        .enumerate()
+        .map(|(idx, rect)| {
+            let (min_x, min_y, max_x, max_y) = minimal_bounding_rect(&rect.points());
+            utils::Bbox {
+                index: idx,
+                x1: min_x,
+                y1: min_y,
+                x2: max_x,
+                y2: max_y,
+            }
+        })
+        .collect();
+    let boxes2_bounding_rects: Vec<utils::Bbox<f64>> = boxes2_rects
+        .iter()
+        .enumerate()
+        .map(|(idx, rect)| {
+            let (min_x, min_y, max_x, max_y) = minimal_bounding_rect(&rect.points());
+            utils::Bbox {
+                index: idx,
+                x1: min_x,
+                y1: min_y,
+                x2: max_x,
+                y2: max_y,
+            }
+        })
+        .collect();
+
+    let box1_rtree: RTree<utils::Bbox<f64>> = RTree::bulk_load(boxes1_bounding_rects);
+    let box2_rtree: RTree<utils::Bbox<f64>> = RTree::bulk_load(boxes2_bounding_rects);
+
+    for (box1, box2) in box1_rtree.intersection_candidates_with_other_tree(&box2_rtree) {
+        let area1 = areas1[box1.index];
+        let area2 = areas2[box2.index];
+        let rect1 = boxes1_rects[box1.index];
+        let rect2 = boxes2_rects[box2.index];
+        let intersection = intersection_area(&rect1, &rect2);
+        let union = area1 + area2 - intersection + utils::EPS;
+        // Calculate the enclosing box (C) coordinates
+        let c_x1 = utils::min(box1.x1, box2.x1);
+        let c_y1 = utils::min(box1.y1, box2.y1);
+        let c_x2 = utils::max(box1.x2, box2.x2);
+        let c_y2 = utils::max(box1.y2, box2.y2);
+        let c_area = (c_x2 - c_x1) * (c_y2 - c_y1);
+        let c_area = c_area.to_f64().unwrap();
+        iou_matrix[[box1.index, box2.index]] = intersection / union - ((c_area - union) / c_area);
+    }
+
+    return iou_matrix;
+}
+
 #[cfg(test)]
 mod tests {
     use ndarray::arr2;
@@ -180,4 +276,12 @@ mod tests {
         assert_eq!(giou_matrix[[1, 1]], 1.2611764705882353);
         assert_eq!(giou_matrix, parallel_giou_matrix);
     }
+
+    #[test]
+    fn test_rotated_giou() {
+        let boxes1 = arr2(&[[5.0, 5.0, 2.0, 2.0, 0.0]]);
+        let boxes2 = arr2(&[[4.0, 4.0, 2.0, 2.0, 0.0]]);
+        let rotated_iou_distance_result = rotated_giou_distance(&boxes1, &boxes2);
+        assert_eq!(rotated_iou_distance_result, arr2(&[[-0.07936507936507936]]));
+    }
 }