From f0e88b1982f8a562cf480bdc5e9263aacda8bc88 Mon Sep 17 00:00:00 2001
From: Linas Kondrackis <linas.ko+dev@skiff.com>
Date: Tue, 28 May 2024 09:52:46 +0300
Subject: [PATCH] Select overlap filtering strategy

---
 supervision/__init__.py                       |  1 +
 .../detection/tools/inference_slicer.py       | 26 ++++++++++++++-----
 supervision/detection/utils.py                | 17 ++++++++++++
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/supervision/__init__.py b/supervision/__init__.py
index 1a46226f..6cd5f9ff 100644
--- a/supervision/__init__.py
+++ b/supervision/__init__.py
@@ -46,6 +46,7 @@ from supervision.detection.tools.json_sink import JSONSink
 from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
 from supervision.detection.tools.smoother import DetectionsSmoother
 from supervision.detection.utils import (
+    OverlapFilter,
     box_iou_batch,
     box_non_max_merge,
     box_non_max_suppression,
diff --git a/supervision/detection/tools/inference_slicer.py b/supervision/detection/tools/inference_slicer.py
index 82551434..3302b139 100644
--- a/supervision/detection/tools/inference_slicer.py
+++ b/supervision/detection/tools/inference_slicer.py
@@ -1,10 +1,11 @@
+import warnings
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Callable, Optional, Tuple
 
 import numpy as np
 
 from supervision.detection.core import Detections
-from supervision.detection.utils import move_boxes, move_masks
+from supervision.detection.utils import OverlapFilter, move_boxes, move_masks
 from supervision.utils.image import crop_image
 
 
@@ -50,8 +51,10 @@ class InferenceSlicer:
             `(width, height)`.
         overlap_ratio_wh (Tuple[float, float]): Overlap ratio between consecutive
             slices in the format `(width_ratio, height_ratio)`.
-        iou_threshold (Optional[float]): Intersection over Union (IoU) threshold
-            used for non-max suppression.
+        overlap_filter (OverlapFilter): Strategy for
+            filtering or merging overlapping detections in slices.
+        iou_threshold (float): Intersection over Union (IoU) threshold
+            used when filtering by overlap.
         callback (Callable): A function that performs inference on a given image
             slice and returns detections.
         thread_workers (int): Number of threads for parallel execution.
@@ -68,12 +71,14 @@ class InferenceSlicer:
         callback: Callable[[np.ndarray], Detections],
         slice_wh: Tuple[int, int] = (320, 320),
         overlap_ratio_wh: Tuple[float, float] = (0.2, 0.2),
-        iou_threshold: Optional[float] = 0.5,
+        overlap_filter: OverlapFilter = OverlapFilter.NON_MAX_SUPPRESSION,
+        iou_threshold: float = 0.5,
         thread_workers: int = 1,
     ):
         self.slice_wh = slice_wh
         self.overlap_ratio_wh = overlap_ratio_wh
         self.iou_threshold = iou_threshold
+        self.overlap_filter = overlap_filter
         self.callback = callback
         self.thread_workers = thread_workers
 
@@ -124,9 +129,16 @@ class InferenceSlicer:
             for future in as_completed(futures):
                 detections_list.append(future.result())
 
-        return Detections.merge(detections_list=detections_list).with_nms(
-            threshold=self.iou_threshold
-        )
+        merged = Detections.merge(detections_list=detections_list)
+        if self.overlap_filter == OverlapFilter.NONE:
+            return merged
+        elif self.overlap_filter == OverlapFilter.NON_MAX_SUPPRESSION:
+            return merged.with_nms(threshold=self.iou_threshold)
+        elif self.overlap_filter == OverlapFilter.NON_MAX_MERGE:
+            return merged.with_nmm(threshold=self.iou_threshold)
+        else:
+            warnings.warn(f"Invalid overlap filter strategy: {self.overlap_filter}")
+            return merged
 
     def _run_callback(self, image, offset) -> Detections:
         """
diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py
index 1ca48791..86b730be 100644
--- a/supervision/detection/utils.py
+++ b/supervision/detection/utils.py
@@ -1,3 +1,4 @@
+from enum import Enum
 from itertools import chain
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -1056,3 +1057,19 @@ def contains_multiple_segments(
         mask_uint8, labels, connectivity=connectivity
     )
     return number_of_labels > 2
+
+
+class OverlapFilter(Enum):
+    """
+    Enum specifying the strategy for filtering overlapping detections.
+
+    Attributes:
+        NONE: Do not filter detections based on overlap.
+        NON_MAX_SUPPRESSION: Filter detections using non-max suppression.
+        NON_MAX_MERGE: Merge detections with non-max-merging instead of
+            discarding them.
+    """
+
+    NONE = "none"
+    NON_MAX_SUPPRESSION = "non_max_suppression"
+    NON_MAX_MERGE = "non_max_merge"