refactor: small optimizations (#19)

Smirkey · Dec 2, 2023 · 03270eb · 03270eb
1 parent 77a2e82
commit 03270eb
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 135 deletions.
diff --git a/bindings/python/powerboxes/__init__.py b/bindings/python/powerboxes/__init__.py
@@ -1,5 +1,5 @@
 import numpy as np
-
+import numpy.typing as npt
 from ._boxes import (
     _dtype_to_func_box_areas,
     _dtype_to_func_box_convert,
@@ -13,13 +13,25 @@
     _dtype_to_func_iou_distance,
     _dtype_to_func_parallel_iou_distance,
 )
-
+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"
+supported_dtypes = [
+    "float64",
+    "float32",
+    "int16",
+    "int32",
+    "int64",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+]
 __version__ = "0.1.3"
 
+T = TypeVar("T", bound=Union[np.float64, np.float32, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64])
 
-def iou_distance(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
+def iou_distance(boxes1: npt.NDArray[T], boxes2: npt.NDArray[T]) -> npt.NDArray[np.float64]:
     """Computes pairwise box iou distances.
 
     Args:
@@ -41,7 +53,7 @@ def iou_distance(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
         raise ValueError(BOXES_NOT_SAME_TYPE)
 
 
-def parallel_iou_distance(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
+def parallel_iou_distance(boxes1: npt.NDArray[T], boxes2: npt.NDArray[T]) -> npt.NDArray[np.float64]:
     """Computes pairwise box iou distances, in parallel.
 
     Args:
@@ -63,7 +75,7 @@ def parallel_iou_distance(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
         raise ValueError(BOXES_NOT_SAME_TYPE)
 
 
-def parallel_giou_distance(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
+def parallel_giou_distance(boxes1: npt.NDArray[T], boxes2: npt.NDArray[T]) -> npt.NDArray[np.float64]:
     """Computes pairwise box giou distances, in parallel.
 
     see: https://giou.stanford.edu/
@@ -87,7 +99,7 @@ def parallel_giou_distance(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray
         raise ValueError(BOXES_NOT_SAME_TYPE)
 
 
-def giou_distance(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
+def giou_distance(boxes1: npt.NDArray[T], boxes2: npt.NDArray[T]) -> npt.NDArray[np.float64]:
     """Computes pairwise box giou distances.
 
     see: https://giou.stanford.edu/
@@ -111,7 +123,7 @@ def giou_distance(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
         raise ValueError(BOXES_NOT_SAME_TYPE)
 
 
-def remove_small_boxes(boxes: np.ndarray, min_size) -> np.ndarray:
+def remove_small_boxes(boxes: npt.NDArray[T], min_size) -> npt.NDArray[T]:
     """Removes boxes with area less than min_area.
 
     Args:
@@ -129,7 +141,7 @@ def remove_small_boxes(boxes: np.ndarray, min_size) -> np.ndarray:
     return _dtype_to_func_remove_small_boxes[boxes.dtype](boxes, min_size)
 
 
-def boxes_areas(boxes: np.ndarray) -> np.ndarray:
+def boxes_areas(boxes: npt.NDArray[T]) -> npt.NDArray[np.float64]:
     """Computes areas of boxes.
 
     Args:
@@ -143,7 +155,7 @@ def boxes_areas(boxes: np.ndarray) -> np.ndarray:
     return _dtype_to_func_box_areas[boxes.dtype](boxes)
 
 
-def box_convert(boxes: np.ndarray, in_fmt: str, out_fmt: str) -> np.ndarray:
+def box_convert(boxes: npt.NDArray[T], in_fmt: str, out_fmt: str) -> npt.NDArray[T]:
     """Converts boxes from one format to another.
 
     Available formats are:
@@ -172,4 +184,5 @@ def box_convert(boxes: np.ndarray, in_fmt: str, out_fmt: str) -> np.ndarray:
     "box_convert",
     "giou_distance",
     "parallel_giou_distance",
+    "supported_dtypes"
 ]
diff --git a/bindings/tests/test_dtypes.py b/bindings/tests/test_dtypes.py
@@ -10,23 +10,13 @@
     parallel_giou_distance,
     parallel_iou_distance,
     remove_small_boxes,
+    supported_dtypes
 )
 
+np.random.seed(42)
 
-@pytest.mark.parametrize(
-    "dtype",
-    [
-        "float64",
-        "float32",
-        "int16",
-        "int32",
-        "int64",
-        "uint8",
-        "uint16",
-        "uint32",
-        "uint64",
-    ],
-)
+
+@pytest.mark.parametrize("dtype", supported_dtypes)
 def test_giou_distance(dtype):
     boxes1 = np.random.random((100, 4))
     boxes2 = np.random.random((100, 4))
@@ -45,20 +35,7 @@ def test_giou_distance_bad_inputs():
         giou_distance("bonjour", "how are you?")
 
 
-@pytest.mark.parametrize(
-    "dtype",
-    [
-        "float64",
-        "float32",
-        "int16",
-        "int32",
-        "int64",
-        "uint8",
-        "uint16",
-        "uint32",
-        "uint64",
-    ],
-)
+@pytest.mark.parametrize("dtype", supported_dtypes)
 def test_parallel_giou_distance(dtype):
     boxes1 = np.random.random((100, 4))
     boxes2 = np.random.random((100, 4))
@@ -77,20 +54,7 @@ def test_parallel_giou_distance_bad_inputs():
         parallel_giou_distance("bonjour", "how are you?")
 
 
-@pytest.mark.parametrize(
-    "dtype",
-    [
-        "float64",
-        "float32",
-        "int16",
-        "int32",
-        "int64",
-        "uint8",
-        "uint16",
-        "uint32",
-        "uint64",
-    ],
-)
+@pytest.mark.parametrize("dtype", supported_dtypes)
 def test_parallel_iou_distance(dtype):
     boxes1 = np.random.random((100, 4))
     boxes2 = np.random.random((100, 4))
@@ -109,20 +73,7 @@ def test_parallel_iou_distance_bad_inputs():
         parallel_iou_distance("bonjour", "how are you?")
 
 
-@pytest.mark.parametrize(
-    "dtype",
-    [
-        "float64",
-        "float32",
-        "int16",
-        "int32",
-        "int64",
-        "uint8",
-        "uint16",
-        "uint32",
-        "uint64",
-    ],
-)
+@pytest.mark.parametrize("dtype", supported_dtypes)
 def test_iou_distance(dtype):
     boxes1 = np.random.random((100, 4))
     boxes2 = np.random.random((100, 4))
@@ -141,20 +92,7 @@ def test_iou_distance_bad_inputs():
         iou_distance("bonjour", "how are you?")
 
 
-@pytest.mark.parametrize(
-    "dtype",
-    [
-        "float64",
-        "float32",
-        "int16",
-        "int32",
-        "int64",
-        "uint8",
-        "uint16",
-        "uint32",
-        "uint64",
-    ],
-)
+@pytest.mark.parametrize("dtype", supported_dtypes)
 def test_remove_small_boxes(dtype):
     boxes = np.random.random((100, 4))
     remove_small_boxes(boxes.astype(dtype), 0.4)
@@ -165,20 +103,7 @@ def test_remove_small_boxes_bad_inputs():
         remove_small_boxes("bonjour", "how are you?")
 
 
-@pytest.mark.parametrize(
-    "dtype",
-    [
-        "float64",
-        "float32",
-        "int16",
-        "int32",
-        "int64",
-        "uint8",
-        "uint16",
-        "uint32",
-        "uint64",
-    ],
-)
+@pytest.mark.parametrize("dtype", supported_dtypes)
 def test_boxes_areas(dtype):
     boxes = np.random.random((100, 4))
     boxes_areas(boxes.astype(dtype))
@@ -189,20 +114,7 @@ def test_boxes_areas_bad_inpus():
         boxes_areas("hey")
 
 
-@pytest.mark.parametrize(
-    "dtype",
-    [
-        "float64",
-        "float32",
-        "int16",
-        "int32",
-        "int64",
-        "uint8",
-        "uint16",
-        "uint32",
-        "uint64",
-    ],
-)
+@pytest.mark.parametrize("dtype", supported_dtypes)
 def test_box_convert(dtype):
     boxes = np.random.random((100, 4))
     box_convert(boxes.astype(dtype), "xyxy", "xywh")

diff --git a/bindings/tests/test_speed.py b/bindings/tests/test_speed.py
@@ -8,80 +8,94 @@
     parallel_giou_distance,
     parallel_iou_distance,
     remove_small_boxes,
+    supported_dtypes
 )
 
+np.random.seed(42)
 
 @pytest.mark.benchmark(group="giou_distance")
-def test_giou_distance(benchmark):
-    boxes1 = np.random.random((100, 4))
-    boxes2 = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_giou_distance(benchmark, dtype):
+    boxes1 = np.random.random((100, 4)).astype(dtype)
+    boxes2 = np.random.random((100, 4)).astype(dtype)
     benchmark(giou_distance, boxes1, boxes2)
 
 
 @pytest.mark.benchmark(group="parallel_giou_distance")
-def test_parallel_giou_distance(benchmark):
-    boxes1 = np.random.random((100, 4))
-    boxes2 = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_parallel_giou_distance(benchmark, dtype):
+    boxes1 = np.random.random((100, 4)).astype(dtype)
+    boxes2 = np.random.random((100, 4)).astype(dtype)
     benchmark(parallel_giou_distance, boxes1, boxes2)
 
 
 @pytest.mark.benchmark(group="iou_distance")
-def test_iou_distance(benchmark):
-    boxes1 = np.random.random((100, 4))
-    boxes2 = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_iou_distance(benchmark, dtype):
+    boxes1 = np.random.random((100, 4)).astype(dtype)
+    boxes2 = np.random.random((100, 4)).astype(dtype)
     benchmark(iou_distance, boxes1, boxes2)
 
 
 @pytest.mark.benchmark(group="parallel_iou_distance")
-def test_parallel_iou_distance(benchmark):
-    boxes1 = np.random.random((100, 4))
-    boxes2 = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_parallel_iou_distance(benchmark, dtype):
+    boxes1 = np.random.random((100, 4)).astype(dtype)
+    boxes2 = np.random.random((100, 4)).astype(dtype)
     benchmark(parallel_iou_distance, boxes1, boxes2)
 
 
 @pytest.mark.benchmark(group="remove_small_boxes")
-def test_remove_small_boxes(benchmark):
-    boxes = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_remove_small_boxes(benchmark, dtype):
+    boxes = np.random.random((100, 4)).astype(dtype)
     benchmark(remove_small_boxes, boxes, 0.4)
 
 
 @pytest.mark.benchmark(group="remove_small_boxes")
-def test_boxes_areas(benchmark):
-    boxes = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_boxes_areas(benchmark, dtype):
+    boxes = np.random.random((100, 4)).astype(dtype)
     benchmark(boxes_areas, boxes)
 
 
 @pytest.mark.benchmark(group="box_convert")
-def test_box_convert_xyxy_xywh(benchmark):
-    boxes = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_box_convert_xyxy_xywh(benchmark, dtype):
+    boxes = np.random.random((100, 4)).astype(dtype)
     benchmark(box_convert, boxes, "xyxy", "xywh")
 
 
 @pytest.mark.benchmark(group="box_convert")
-def test_box_convert_xyxy_cxcywh(benchmark):
-    boxes = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_box_convert_xyxy_cxcywh(benchmark, dtype):
+    boxes = np.random.random((100, 4)).astype(dtype)
     benchmark(box_convert, boxes, "xyxy", "cxcywh")
 
 
 @pytest.mark.benchmark(group="box_convert")
-def test_box_convert_cxcywh_xywh(benchmark):
-    boxes = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_box_convert_cxcywh_xywh(benchmark, dtype):
+    boxes = np.random.random((100, 4)).astype(dtype)
     benchmark(box_convert, boxes, "cxcywh", "xywh")
 
 
 @pytest.mark.benchmark(group="box_convert")
-def test_box_convert_cxcywh_xyxy(benchmark):
-    boxes = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_box_convert_cxcywh_xyxy(benchmark, dtype):
+    boxes = np.random.random((100, 4)).astype(dtype)
     benchmark(box_convert, boxes, "cxcywh", "xywh")
 
 
 @pytest.mark.benchmark(group="box_convert")
-def test_box_convert_xywh_cxcywh(benchmark):
-    boxes = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_box_convert_xywh_cxcywh(benchmark, dtype):
+    boxes = np.random.random((100, 4)).astype(dtype)
     benchmark(box_convert, boxes, "xywh", "cxcywh")
 
 
 @pytest.mark.benchmark(group="box_convert")
-def test_box_convert_xywh_xyxy(benchmark):
-    boxes = np.random.random((100, 4))
+@pytest.mark.parametrize("dtype", supported_dtypes)
+def test_box_convert_xywh_xyxy(benchmark, dtype):
+    boxes = np.random.random((100, 4)).astype(dtype)
     benchmark(box_convert, boxes, "xywh", "xyxy")