Merge pull request #938 from roboflow/mask-iou-batch-memory-issue

Limit memory usage for mask_iou_batch

supervision/detection/utils.py

@@ -59,8 +59,11 @@ def box_area(box):
     return area_inter / (area_true[:, None] + area_detection - area_inter)
 
 
-def mask_iou_batch(masks_true: np.ndarray, masks_detection: np.ndarray) -> np.ndarray:
+def _mask_iou_batch_split(
+    masks_true: np.ndarray, masks_detection: np.ndarray
+) -> np.ndarray:
     """
+    Internal function.
     Compute Intersection over Union (IoU) of two sets of masks -
         `masks_true` and `masks_detection`.
 
@@ -74,9 +77,9 @@ def mask_iou_batch(masks_true: np.ndarray, masks_detection: np.ndarray) -> np.nd
     intersection_area = np.logical_and(masks_true[:, None], masks_detection).sum(
         axis=(2, 3)
     )
+
     masks_true_area = masks_true.sum(axis=(1, 2))
     masks_detection_area = masks_detection.sum(axis=(1, 2))
-
     union_area = masks_true_area[:, None] + masks_detection_area - intersection_area
 
     return np.divide(
@@ -87,6 +90,48 @@ def mask_iou_batch(masks_true: np.ndarray, masks_detection: np.ndarray) -> np.nd
     )
 
 
+def mask_iou_batch(
+    masks_true: np.ndarray,
+    masks_detection: np.ndarray,
+    memory_limit: int = 1024 * 1024 * 5,
+) -> np.ndarray:
+    """
+    Compute Intersection over Union (IoU) of two sets of masks -
+        `masks_true` and `masks_detection`.
+
+    Args:
+        masks_true (np.ndarray): 3D `np.ndarray` representing ground-truth masks.
+        masks_detection (np.ndarray): 3D `np.ndarray` representing detection masks.
+        memory_limit (int, optional): memory limit in bytes, default is 5GB.
+
+    Returns:
+        np.ndarray: Pairwise IoU of masks from `masks_true` and `masks_detection`.
+    """
+    memory = (
+        masks_true.shape[0]
+        * masks_true.shape[1]
+        * masks_true.shape[2]
+        * masks_detection.shape[0]
+    )
+    if memory <= memory_limit:  # 5GB memory
+        return _mask_iou_batch_split(masks_true, masks_detection)
+
+    ious = []
+    step = max(
+        memory_limit
+        // (
+            masks_detection.shape[0]
+            * masks_detection.shape[1]
+            * masks_detection.shape[2]
+        ),
+        1,
+    )
+    for i in range(0, masks_true.shape[0], step):
+        ious.append(_mask_iou_batch_split(masks_true[i : i + step], masks_detection))
+
+    return np.vstack(ious)
+
+
 def resize_masks(masks: np.ndarray, max_dimension: int = 640) -> np.ndarray:
     """
     Resize all masks in the array to have a maximum dimension of max_dimension,