Repository: Project-MONAI/MONAI

Issue:
`GridPatch` Performance Issue
**Describe the bug**
`GridPatch` and `RandGridPatch` are slower than their similar component `TileOnGrid`.

**To Reproduce**
```python
import numpy as np
from monai.transforms import GridPatchd
from monai.apps.pathology.transforms import TileOnGridd
from numpy.testing import assert_array_equal, assert_equal

# input configs
patch_size = 256
num_patches = 90
h, w = 10000, 1000
img = {"image": np.random.randint(0, 255, (3, h, w))}

# create components objects
patcher = GridPatchd(
    keys=["image"],
    patch_size=(patch_size, patch_size),
    num_patches=num_patches,
    sort_fn="min",
    threshold=0.999 * 3 * 255 * patch_size * patch_size,
    pad_mode=None,
    constant_values=255,
)
tiler = TileOnGridd(
    keys=["image"],
    tile_count=num_patches,
    tile_size=patch_size,
    background_val=255,
    return_list_of_dicts=True,
)

# Make sure the outputs are identical
patch_output = patcher(img)
tile_output = tiler(img)
assert_equal(len(patch_output), len(tile_output))
for p, t in zip(patch_output, tile_output):
    assert_array_equal(p["image"], t["image"])
```
```python
%%timeit -n 10 -r 10
patch_output = patcher(img)
```
```
208 ms ± 16.4 ms per loop (mean ± std. dev. of 10 runs, 10 loops each)
```
```python
%%timeit -n 10 -r 10
tile_output = tiler(img)
```
```
104 ms ± 2.38 ms per loop (mean ± std. dev. of 10 runs, 10 loops each)
```


Discussion:

Resolving patch:
diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py
--- a/monai/transforms/spatial/array.py
+++ b/monai/transforms/spatial/array.py
@@ -2633,11 +2633,8 @@ class GridPatch(Transform):
         num_patches: number of patches to return. Defaults to None, which returns all the available patches.
         overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0).
             If only one float number is given, it will be applied to all dimensions. Defaults to 0.0.
-        sort_fn: a callable or string that defines the order of the patches to be returned. If it is a callable, it
-            will be passed directly to the `key` argument of `sorted` function. The string can be "min" or "max",
-            which are, respectively, the minimum and maximum of the sum of intensities of a patch across all dimensions
-            and channels. Also "random" creates a random order of patches.
-            By default no sorting is being done and patches are returned in a row-major order.
+        sort_fn: when `num_patches` is provided, it determines if keep patches with highest values (`"max"`),
+            lowest values (`"min"`), or in their default order (`None`). Default to None.
         threshold: a value to keep only the patches whose sum of intensities are less than the threshold.
             Defaults to no filtering.
         pad_mode: refer to NumpyPadMode and PytorchPadMode. If None, no padding will be applied. Defaults to ``"constant"``.
@@ -2653,7 +2650,7 @@ def __init__(
         offset: Optional[Sequence[int]] = None,
         num_patches: Optional[int] = None,
         overlap: Union[Sequence[float], float] = 0.0,
-        sort_fn: Optional[Union[Callable, str]] = None,
+        sort_fn: Optional[str] = None,
         threshold: Optional[float] = None,
         pad_mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.CONSTANT,
         **pad_kwargs,
@@ -2664,24 +2661,45 @@ def __init__(
         self.pad_kwargs = pad_kwargs
         self.overlap = overlap
         self.num_patches = num_patches
-        self.sort_fn: Optional[Callable]
-        if isinstance(sort_fn, str):
-            self.sort_fn = GridPatchSort.get_sort_fn(sort_fn)
-        else:
-            self.sort_fn = sort_fn
-
+        self.sort_fn = sort_fn.lower() if sort_fn else None
         self.threshold = threshold
-        if threshold:
-            self.filter_fn = self.threshold_fn
-        else:
-            self.filter_fn = self.one_fn
 
-    @staticmethod
-    def one_fn(patch):
-        return True
+    def filter_threshold(self, image_np: np.ndarray, locations: np.ndarray):
+        """
+        Filter the patches and their locations according to a threshold
+        Args:
+            image: a numpy.ndarray representing a stack of patches
+            location: a numpy.ndarray representing the stack of location of each patch
+        """
+        if self.threshold is not None:
+            n_dims = len(image_np.shape)
+            idx = np.argwhere(image_np.sum(axis=tuple(range(1, n_dims))) < self.threshold).reshape(-1)
+            image_np = image_np[idx]
+            locations = locations[idx]
+        return image_np, locations
 
-    def threshold_fn(self, patch):
-        return patch.sum() < self.threshold
+    def filter_count(self, image_np: np.ndarray, locations: np.ndarray):
+        """
+        Sort the patches based on the sum of their intensity, and just keep `self.num_patches` of them.
+        Args:
+            image: a numpy.ndarray representing a stack of patches
+            location: a numpy.ndarray representing the stack of location of each patch
+        """
+        if self.sort_fn is None:
+            image_np = image_np[: self.num_patches]
+            locations = locations[: self.num_patches]
+        elif self.num_patches is not None:
+            n_dims = len(image_np.shape)
+            if self.sort_fn == GridPatchSort.MIN:
+                idx = np.argsort(image_np.sum(axis=tuple(range(1, n_dims))))
+            elif self.sort_fn == GridPatchSort.MAX:
+                idx = np.argsort(-image_np.sum(axis=tuple(range(1, n_dims))))
+            else:
+                raise ValueError(f'`sort_fn` should be either "min", "max" or None! {self.sort_fn} provided!')
+            idx = idx[: self.num_patches]
+            image_np = image_np[idx]
+            locations = locations[idx]
+        return image_np, locations
 
     def __call__(self, array: NdarrayOrTensor):
         # create the patch iterator which sweeps the image row-by-row
@@ -2695,25 +2713,28 @@ def __call__(self, array: NdarrayOrTensor):
             mode=self.pad_mode,
             **self.pad_kwargs,
         )
+        patches = list(zip(*patch_iterator))
+        patched_image = np.array(patches[0])
+        locations = np.array(patches[1])[:, 1:, 0]  # only keep the starting location
 
-        if self.sort_fn is not None:
-            patch_iterator = sorted(patch_iterator, key=self.sort_fn)
+        # Filter patches
+        if self.num_patches:
+            patched_image, locations = self.filter_count(patched_image, locations)
+        elif self.threshold:
+            patched_image, locations = self.filter_threshold(patched_image, locations)
 
-        output = [
-            (convert_to_dst_type(src=patch, dst=array)[0], convert_to_dst_type(src=slices[..., 0], dst=array)[0])
-            for patch, slices in patch_iterator
-            if self.filter_fn(patch)
-        ]
+        # Convert to original data type
+        output = list(
+            zip(convert_to_dst_type(src=patched_image, dst=array)[0], convert_to_dst_type(src=locations, dst=array)[0])
+        )
 
-        if self.num_patches:
-            output = output[: self.num_patches]
-            if len(output) < self.num_patches:
-                patch = convert_to_dst_type(
-                    src=np.full((array.shape[0], *self.patch_size), self.pad_kwargs.get("constant_values", 0)),
-                    dst=array,
-                )[0]
-                start_location = convert_to_dst_type(src=np.zeros((len(self.patch_size), 1)), dst=array)[0]
-                output += [(patch, start_location)] * (self.num_patches - len(output))
+        # Pad the patch list to have the requested number of patches
+        if self.num_patches and len(output) < self.num_patches:
+            patch = convert_to_dst_type(
+                src=np.full((array.shape[0], *self.patch_size), self.pad_kwargs.get("constant_values", 0)), dst=array
+            )[0]
+            start_location = convert_to_dst_type(src=np.zeros((len(self.patch_size), 1)), dst=array)[0]
+            output += [(patch, start_location)] * (self.num_patches - len(output))
 
         return output
 
@@ -2732,11 +2753,8 @@ class RandGridPatch(GridPatch, RandomizableTransform):
         num_patches: number of patches to return. Defaults to None, which returns all the available patches.
         overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0).
             If only one float number is given, it will be applied to all dimensions. Defaults to 0.0.
-        sort_fn: a callable or string that defines the order of the patches to be returned. If it is a callable, it
-            will be passed directly to the `key` argument of `sorted` function. The string can be "min" or "max",
-            which are, respectively, the minimum and maximum of the sum of intensities of a patch across all dimensions
-            and channels. Also "random" creates a random order of patches.
-            By default no sorting is being done and patches are returned in a row-major order.
+        sort_fn: when `num_patches` is provided, it determines if keep patches with highest values (`"max"`),
+            lowest values (`"min"`), or in their default order (`None`). Default to None.
         threshold: a value to keep only the patches whose sum of intensities are less than the threshold.
             Defaults to no filtering.
         pad_mode: refer to NumpyPadMode and PytorchPadMode. If None, no padding will be applied. Defaults to ``"constant"``.
@@ -2753,7 +2771,7 @@ def __init__(
         max_offset: Optional[Union[Sequence[int], int]] = None,
         num_patches: Optional[int] = None,
         overlap: Union[Sequence[float], float] = 0.0,
-        sort_fn: Optional[Union[Callable, str]] = None,
+        sort_fn: Optional[str] = None,
         threshold: Optional[float] = None,
         pad_mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.CONSTANT,
         **pad_kwargs,
diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py
--- a/monai/transforms/spatial/dictionary.py
+++ b/monai/transforms/spatial/dictionary.py
@@ -17,7 +17,7 @@
 
 from copy import deepcopy
 from enum import Enum
-from typing import Any, Callable, Dict, Hashable, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Hashable, List, Mapping, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
@@ -2216,11 +2216,8 @@ class GridPatchd(MapTransform):
             np.random.randint(0, patch_size, 2) creates random start between 0 and `patch_size` for a 2D image.
         num_patches: number of patches to return. Defaults to None, which returns all the available patches.
         overlap: amount of overlap between patches in each dimension. Default to 0.0.
-        sort_fn: a callable or string that defines the order of the patches to be returned. If it is a callable, it
-            will be passed directly to the `key` argument of `sorted` function. The string can be "min" or "max",
-            which are, respectively, the minimum and maximum of the sum of intensities of a patch across all dimensions
-            and channels. Also "random" creates a random order of patches.
-            By default no sorting is being done and patches are returned in a row-major order.
+        sort_fn: when `num_patches` is provided, it determines if keep patches with highest values (`"max"`),
+            lowest values (`"min"`), or in their default order (`None`). Default to None.
         threshold: a value to keep only the patches whose sum of intensities are less than the threshold.
             Defaults to no filtering.
         pad_mode: refer to NumpyPadMode and PytorchPadMode. If None, no padding will be applied. Defaults to ``"constant"``.
@@ -2246,7 +2243,7 @@ def __init__(
         offset: Optional[Sequence[int]] = None,
         num_patches: Optional[int] = None,
         overlap: float = 0.0,
-        sort_fn: Optional[Union[Callable, str]] = None,
+        sort_fn: Optional[str] = None,
         threshold: Optional[float] = None,
         pad_mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.CONSTANT,
         allow_missing_keys: bool = False,
@@ -2300,11 +2297,8 @@ class RandGridPatchd(RandomizableTransform, MapTransform):
         num_patches: number of patches to return. Defaults to None, which returns all the available patches.
         overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0).
             If only one float number is given, it will be applied to all dimensions. Defaults to 0.0.
-        sort_fn: a callable or string that defines the order of the patches to be returned. If it is a callable, it
-            will be passed directly to the `key` argument of `sorted` function. The string can be "min" or "max",
-            which are, respectively, the minimum and maximum of the sum of intensities of a patch across all dimensions
-            and channels. Also "random" creates a random order of patches.
-            By default no sorting is being done and patches are returned in a row-major order.
+        sort_fn: when `num_patches` is provided, it determines if keep patches with highest values (`"max"`),
+            lowest values (`"min"`), or in their default order (`None`). Default to None.
         threshold: a value to keep only the patches whose sum of intensities are less than the threshold.
             Defaults to no filtering.
         pad_mode: refer to NumpyPadMode and PytorchPadMode. If None, no padding will be applied. Defaults to ``"constant"``.
@@ -2332,7 +2326,7 @@ def __init__(
         max_offset: Optional[Union[Sequence[int], int]] = None,
         num_patches: Optional[int] = None,
         overlap: float = 0.0,
-        sort_fn: Optional[Union[Callable, str]] = None,
+        sort_fn: Optional[str] = None,
         threshold: Optional[float] = None,
         pad_mode: Union[NumpyPadMode, PytorchPadMode, str] = NumpyPadMode.CONSTANT,
         allow_missing_keys: bool = False,


Test patch: