feat: add simple patching func

Author: okunator

cellseg_models_pytorch/datasets/dataset_writers/_base_writer.py

@@ -6,7 +6,7 @@
 from pathos.multiprocessing import ThreadPool as Pool
 
 from ...transforms.albu_transforms import IMG_TRANSFORMS, compose
-from ...utils import FileHandler, TilerStitcher, fix_duplicates, remap_label
+from ...utils import FileHandler, fix_duplicates, get_patches, remap_label
 
 __all__ = ["BaseWriter"]
 
@@ -184,11 +184,7 @@ def _get_tiles(
         masks: Union[Dict[str, np.ndarray], None] = None,
     ) -> Tuple[Dict[str, np.ndarray], Union[Dict[str, np.ndarray], None]]:
         """Do tiling on an image and corresponding masks if there are such."""
-        # Init Tilers
-        im_tiler = TilerStitcher(
-            im_shape=im.shape, patch_shape=self.patch_size + (3,), stride=self.stride
-        )
-        im_tiles = im_tiler.patch(im)
+        im_tiles = get_patches(im, self.stride, self.patch_size)[0]
 
         # Tile masks if there are such.
         mask_tiles = None
@@ -204,18 +200,18 @@ def _get_tiles(
             if "sem_map" in masks.keys():
                 sem = masks["sem_map"]
 
-            mask_tiler = TilerStitcher(
-                im_shape=inst.shape,
-                patch_shape=self.patch_size + (1,),
-                stride=self.stride,
-            )
-
             if inst is not None:
-                mask_tiles["inst_map"] = mask_tiler.patch(inst).squeeze()
+                mask_tiles["inst_map"] = get_patches(
+                    inst, self.stride, self.patch_size
+                )[0]
             if types is not None:
-                mask_tiles["type_map"] = mask_tiler.patch(types).squeeze()
+                mask_tiles["type_map"] = get_patches(
+                    types, self.stride, self.patch_size
+                )[0]
             if sem is not None:
-                mask_tiles["sem_map"] = mask_tiler.patch(sem).squeeze()
+                mask_tiles["sem_map"] = get_patches(sem, self.stride, self.patch_size)[
+                    0
+                ]
 
         return im_tiles, mask_tiles
 

cellseg_models_pytorch/utils/patching.py

@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import Dict, Tuple
 
 import numpy as np
 import torch
@@ -10,6 +10,7 @@
     "extract_patches_numpy",
     "stitch_patches_numpy",
     "TilerStitcher",
+    "get_patches",
 ]
 
 
@@ -246,6 +247,143 @@ def stitch_patches_torch(
     return output
 
 
+def _get_margins_and_pad(
+    first_endpoint: int, img_size: int, stride: int, pad: int = None
+) -> Tuple[int, int]:
+    """Get the number of slices needed for one direction and the overlap."""
+    pad = int(pad) if pad is not None else 20  # at least some padding needed
+    img_size += pad
+
+    n = 1
+    mod = 0
+    end = first_endpoint
+    while True:
+        n += 1
+        end += stride
+
+        if end > img_size:
+            mod = end - img_size
+            break
+        elif end == img_size:
+            break
+
+    return n, mod + pad
+
+
+def _get_slices(
+    stride: int,
+    patch_size: Tuple[int, int],
+    img_size: Tuple[int, int],
+    pad: int = None,
+) -> Tuple[Dict[str, slice], int, int]:
+    """Get all the overlapping slices in a dictionary and the needed paddings."""
+    y_end, x_end = patch_size
+    nrows, pady = _get_margins_and_pad(y_end, img_size[0], stride, pad=pad)
+    ncols, padx = _get_margins_and_pad(x_end, img_size[1], stride, pad=pad)
+
+    xyslices = {}
+    for row in range(nrows):
+        for col in range(ncols):
+            y_start = row * stride
+            y_end = y_start + patch_size[0]
+            x_start = col * stride
+            x_end = x_start + patch_size[1]
+            xyslices[f"y-{y_start}_x-{x_start}"] = (
+                slice(y_start, y_end),
+                slice(x_start, x_end),
+            )
+
+    return xyslices, pady, padx, nrows, ncols
+
+
+def get_patches(
+    arr: np.ndarray, stride: int, patch_size: Tuple[int, int], padding: int = None
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Tuple[int, ...], int, int]:
+    """Patch an input array to overlapping or non-overlapping patches.
+
+    NOTE: some padding is applied by default to make the arr divisible by patch_size.
+
+    Parameters
+    ----------
+        arr : np.ndarray
+            An array of shape: (H, W, C) or (H, W).
+        stride : int
+            Stride of the sliding window.
+        patch_size : Tuple[int, int]
+            Height and width of the patch
+        padding : int, optional
+            Size of reflection padding.
+
+    Returns
+    --------
+        Tuple[np.ndarray, np.ndarray, np.ndarray, Tuple[int, ...], int, int]:
+            - Batched input array of shape: (n_patches, ph, pw)|(n_patches, ph, pw, C)
+            - Batched repeats array of shape: (n_patches, ph, pw) int32
+            - Repeat matrix of shape (H + pady, W + padx). Dtype: int32
+            - The shape of the padded input array.
+            - nrows
+            - ncols
+    """
+    shape = arr.shape
+    if len(shape) == 2:
+        arr_type = "HW"
+    elif len(shape) == 3:
+        arr_type = "HWC"
+    else:
+        raise ValueError("`arr` needs to be either 'HW' or 'HWC' shape.")
+
+    slices, pady, padx, nrows, ncols = _get_slices(
+        stride, patch_size, (shape[0], shape[1]), padding
+    )
+
+    padx, modx = divmod(padx, 2)
+    pady, mody = divmod(pady, 2)
+    padx += modx
+    pady += mody
+
+    pad = [(pady, pady), (padx, padx)]
+    if arr_type == "HWC":
+        pad.append((0, 0))
+
+    arr = np.pad(arr, pad, mode="reflect")
+
+    # init repeats matrix + add padding repeats
+    if padding != 0 or padding is None:
+        repeats = np.ones(arr.shape[:2])
+        repeats[pady:-pady, padx:-padx] = 0
+
+        # corner pads
+        repeats[:pady, :padx] += 1
+        repeats[-pady:, -padx:] += 1
+        repeats[-pady:, :padx] += 1
+        repeats[:pady, -padx:] += 1
+    else:
+        repeats = np.zeros(arr.shape[:2])
+
+    patches = []
+    rep_patches = []
+    for yslice, xslice in slices.values():
+        if arr_type == "HW":
+            patch = arr[yslice, xslice]
+        elif arr_type == "HWC":
+            patch = arr[yslice, xslice, ...]
+
+        rep_patch = repeats[yslice, xslice]
+        repeats[yslice, xslice] += 1
+
+        patches.append(patch)
+        rep_patches.append(rep_patch)
+
+    return (
+        np.array(patches, dtype="uint8"),
+        np.array(rep_patches, dtype="int32"),
+        repeats.astype("int32"),
+        arr.shape,
+        nrows,
+        ncols,
+    )
+
+
 class TilerStitcher:
     def __init__(
         self,

cellseg_models_pytorch/utils/tests/test_patching.py

@@ -7,6 +7,7 @@
     TilerStitcherTorch,
     extract_patches_numpy,
     extract_patches_torch,
+    get_patches,
     stitch_patches_numpy,
     stitch_patches_torch,
 )
@@ -102,3 +103,16 @@ def test_tilerstitchertorch(rand_tensor, padding):
     stitched = ts.backstitch(patches)
 
     assert stitched.shape == rand_tensor.shape
+
+
+@pytest.mark.parametrize("padding", [10, None])
+def test_dict_patching(img_sample, padding):
+    stride = 32
+    patch_size = (64, 64)
+    padding = None
+    patches, _, repeats, padded_shape, _, _ = get_patches(
+        img_sample, stride, patch_size, padding
+    )
+
+    assert patches.shape[1:-1] == patch_size
+    assert repeats.shape == padded_shape[:-1]