Feat/inference slicer segmentation

supervision/detection/tools/inference_slicer.py

@@ -4,20 +4,28 @@
 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: np.ndarray
+) -> 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 (np.ndarray): An array of shape `(2,)` or `(3,)`, size of the image
+            is in format `[width, height]`.
     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 +134,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,46 @@ 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[np.ndarray] = 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`: array of ints, shaped `(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 (np.ndarray, optional): Final shape of the mask in the format
+            `(height, width, ...)`. If provided, the masks will be padded to match this
+            shape. 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_x, size_y = masks.shape[1:] + offset[::-1]
+    if desired_shape is not None:
+        size_y, size_x = desired_shape[:2]
+
+    mask_arr = np.full((masks.shape[0], size_y, size_x), False)
+    mask_arr[
+        :,
+        offset[1] : masks.shape[1] + offset[1],
+        offset[0] : masks.shape[2] + offset[0],
+    ] = masks
+
+    return mask_arr
+
+
 def scale_boxes(xyxy: np.ndarray, factor: float) -> np.ndarray:
     """
     Scale the dimensions of bounding boxes.