Feat/inference slicer segmentation

docs/detection/utils.md

@@ -65,6 +65,12 @@ status: new
 
 :::supervision.detection.utils.move_boxes
 
+<div class="md-typeset">
+  <h2><a href="#supervision.detection.utils.move_masks">move_masks</a></h2>
+</div>
+
+:::supervision.detection.utils.move_masks
+
 <div class="md-typeset">
   <h2><a href="#supervision.detection.utils.scale_boxes">scale_boxes</a></h2>
 </div>

supervision/__init__.py

@@ -52,6 +52,7 @@
     mask_to_polygons,
     mask_to_xyxy,
     move_boxes,
+    move_masks,
     polygon_to_mask,
     polygon_to_xyxy,
     scale_boxes,

supervision/detection/tools/inference_slicer.py

@@ -1,23 +1,35 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Callable, Optional, Tuple
+from typing import Callable, Optional, Tuple, Union
 
 import numpy as np
 
 from supervision.detection.core import Detections
-from supervision.detection.utils import move_boxes
+from supervision.detection.utils import move_boxes, move_masks
 from supervision.utils.image import crop_image
 
 
-def move_detections(detections: Detections, offset: np.array) -> Detections:
+def move_detections(
+    detections: Detections,
+    offset: np.ndarray,
+    image_shape: Optional[Union[Tuple[int, int, int], Tuple[int, int]]] = None,
+) -> Detections:
     """
     Args:
         detections (sv.Detections): Detections object to be moved.
-        offset (np.array): An array of shape `(2,)` containing offset values in format
+        offset (np.ndarray): An array of shape `(2,)` containing offset values in format
             is `[dx, dy]`.
+        image_size (Tuple, optional): A tuple of image shape. Can be  `(2,)` or `(3,)`.
+            Important when moving for segmentation detections, as it defines mask array
+            size.
+
     Returns:
         (sv.Detections) repositioned Detections object.
     """
     detections.xyxy = move_boxes(xyxy=detections.xyxy, offset=offset)
+    if detections.mask is not None:
+        detections.mask = move_masks(
+            masks=detections.mask, offset=offset, desired_shape=image_shape
+        )
     return detections
 
 
@@ -126,7 +138,9 @@ def _run_callback(self, image, offset) -> Detections:
         """
         image_slice = crop_image(image=image, xyxy=offset)
         detections = self.callback(image_slice)
-        detections = move_detections(detections=detections, offset=offset[:2])
+        detections = move_detections(
+            detections=detections, offset=offset[:2], image_shape=image.shape
+        )
 
         return detections
 

supervision/detection/utils.py

@@ -592,6 +592,49 @@ def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray:
     return xyxy + np.hstack([offset, offset])
 
 
+def move_masks(
+    masks: np.ndarray,
+    offset: np.ndarray,
+    desired_shape: Optional[Union[Tuple[int, int, int], Tuple[int, int]]] = None,
+) -> np.ndarray:
+    """
+    Offset the masks in an array by the specified (x, y) amount.
+
+    Note the axis orders:
+
+    - `masks`: array of shape `(n, y, x)`
+    - `offset`: array of ints: `(x, y)`
+    - `desired_shape`: tuple of ints, shaped `(y, x)` or `(y, x, ...)`
+
+    Args:
+        masks (np.ndarray): array of bools
+        offset (np.ndarray): An array of shape `(2,)` containing non-negative int values
+            `[dx, dy]`.
+        desired_shape (Tuple, optional): Final shape of the mask in the format
+            `(height, width)`, `(height, width, ...)`. The masks will be padded to match
+            the first 2 shape dimensions. Note the axis order (y,x)!
+
+    Returns:
+        (np.ndarray) repositioned masks, optionally padded to the specified shape.
+    """
+
+    if offset[0] < 0 or offset[1] < 0:
+        raise ValueError(f"Offset values must be non-negative integers. Got: {offset}")
+
+    size_y, size_x = masks.shape[1:] + offset[::-1]
+    if desired_shape is not None:
+        size_y, size_x = desired_shape[:2]
+
+    mask_array = np.full((masks.shape[0], size_y, size_x), False)
+    mask_array[
+        :,
+        offset[1] : masks.shape[1] + offset[1],
+        offset[0] : masks.shape[2] + offset[0],
+    ] = masks
+
+    return mask_array
+
+
 def scale_boxes(xyxy: np.ndarray, factor: float) -> np.ndarray:
     """
     Scale the dimensions of bounding boxes.