Improve rank filtering performance by removing use of footprint kernel when possible

python/cucim/src/cucim/skimage/_vendored/_ndimage_filters.py

@@ -848,7 +848,8 @@ def _min_or_max_filter(input, size, ftprnt, structure, output, mode, cval,
             input, sizes, output, mode, cval, origin)
 
     origins, int_type = _filters_core._check_nd_args(input, ftprnt,
-                                                     mode, origin, 'footprint')
+                                                     mode, origin, 'footprint',
+                                                     sizes=sizes)
     if structure is not None and structure.ndim != input.ndim:
         raise RuntimeError('structure array has incorrect shape')
 
@@ -1081,28 +1082,51 @@ def get_rank(fs):
 
 def _rank_filter(input, get_rank, size=None, footprint=None, output=None,
                  mode="reflect", cval=0.0, origin=0):
-    _, footprint, _ = _filters_core._check_size_footprint_structure(
-        input.ndim, size, footprint, None, force_footprint=True)
+    sizes, footprint, _ = _filters_core._check_size_footprint_structure(
+        input.ndim, size, footprint, None, force_footprint=False)
     if cval is cupy.nan:
         raise NotImplementedError("NaN cval is unsupported")
     origins, int_type = _filters_core._check_nd_args(input, footprint,
-                                                     mode, origin, 'footprint')
-    if footprint.size == 0:
+                                                     mode, origin, 'footprint',
+                                                     sizes=sizes)
+    has_weights = True
+    if sizes is not None:
+        has_weights = False
+        filter_size = internal.prod(sizes)
+        if filter_size == 0:
+            return cupy.zeros_like(input)
+        footprint_shape = tuple(sizes)
+    elif footprint.size == 0:
         return cupy.zeros_like(input)
-    filter_size = int(footprint.sum())
+    else:
+        footprint_shape = footprint.shape
+        filter_size = int(footprint.sum())
+        if filter_size == footprint.size:
+            # can omit passing the footprint if it is all ones
+            sizes = footprint.shape
+            has_weights = False
+
     rank = get_rank(filter_size)
     if rank < 0 or rank >= filter_size:
         raise RuntimeError('rank not within filter footprint size')
     if rank == 0:
-        return _min_or_max_filter(input, None, footprint, None, output, mode,
-                                  cval, origins, 'min')
-    if rank == filter_size - 1:
-        return _min_or_max_filter(input, None, footprint, None, output, mode,
-                                  cval, origins, 'max')
-    offsets = _filters_core._origins_to_offsets(origins, footprint.shape)
-    kernel = _get_rank_kernel(filter_size, rank, mode, footprint.shape,
-                              offsets, float(cval), int_type)
-    return _filters_core._call_kernel(kernel, input, footprint, output,
+        min_max_op = 'min'
+    elif rank == filter_size - 1:
+        min_max_op = 'max'
+    else:
+        min_max_op = None
+    if min_max_op is not None:
+        if sizes is not None:
+            return _min_or_max_filter(input, sizes[0], None, None, output,
+                                      mode, cval, origins, min_max_op)
+        else:
+            return _min_or_max_filter(input, None, footprint, None, output,
+                                      mode, cval, origins, min_max_op)
+    offsets = _filters_core._origins_to_offsets(origins, footprint_shape)
+    kernel = _get_rank_kernel(filter_size, rank, mode, footprint_shape,
+                              offsets, float(cval), int_type,
+                              has_weights=has_weights)
+    return _filters_core._call_kernel(kernel, input, None, output,
                                       weights_dtype=bool)
 
 
@@ -1134,7 +1158,7 @@ def _get_shell_gap(filter_size):
 
 @cupy._util.memoize(for_each_device=True)
 def _get_rank_kernel(filter_size, rank, mode, w_shape, offsets, cval,
-                     int_type):
+                     int_type, has_weights):
     s_rank = min(rank, filter_size - rank - 1)
     # The threshold was set based on the measurements on a V100
     # TODO(leofang, anaruse): Use Optuna to automatically tune the threshold,
@@ -1184,4 +1208,5 @@ def _get_rank_kernel(filter_size, rank, mode, w_shape, offsets, cval,
         'rank_{}_{}'.format(filter_size, rank),
         'int iv = 0;\nX values[{}];'.format(array_size),
         'values[iv++] = {value};' + found_post, post,
-        mode, w_shape, int_type, offsets, cval, preamble=sorter)
+        mode, w_shape, int_type, offsets, cval, has_weights=has_weights,
+        preamble=sorter)

python/cucim/src/cucim/skimage/_vendored/_ndimage_filters_core.py

@@ -56,14 +56,24 @@ def _convert_1d_args(ndim, weights, origin, axis):
     return weights, tuple(origins)
 
 
-def _check_nd_args(input, weights, mode, origin, wghts_name='filter weights'):
+def _check_nd_args(input, weights, mode, origin, wghts_name='filter weights',
+                   sizes=None):
     _util._check_mode(mode)
-    # Weights must always be less than 2 GiB
-    if weights.nbytes >= (1 << 31):
-        raise RuntimeError('weights must be 2 GiB or less, use FFTs instead')
-    weight_dims = [x for x in weights.shape if x != 0]
-    if len(weight_dims) != input.ndim:
-        raise RuntimeError('{} array has incorrect shape'.format(wghts_name))
+    if weights is not None:
+        # Weights must always be less than 2 GiB
+        if weights.nbytes >= (1 << 31):
+            raise RuntimeError(
+                "weights must be 2 GiB or less, use FFTs instead"
+            )
+        weight_dims = [x for x in weights.shape if x != 0]
+        if len(weight_dims) != input.ndim:
+            raise RuntimeError(
+                "{} array has incorrect shape".format(wghts_name)
+            )
+    elif sizes is None:
+        raise ValueError("must specify either weights array or sizes")
+    else:
+        weight_dims = sizes
     origins = _util._fix_sequence_arg(origin, len(weight_dims), 'origin', int)
     for origin, width in zip(origins, weight_dims):
         _util._check_origin(origin, width)

python/cucim/src/cucim/skimage/filters/_median.py

@@ -29,12 +29,16 @@ def median(image, footprint=None, out=None, mode='nearest', cval=0.0,
     ----------
     image : array-like
         Input image.
-    footprint : ndarray, optional
-        If ``behavior=='rank'``, ``footprint`` is a 2-D array of 1's and 0's.
-        If ``behavior=='ndimage'``, ``footprint`` is a N-D array of 1's and 0's
-        with the same number of dimension than ``image``.
-        If None, ``footprint`` will be a N-D array with 3 elements for each
-        dimension (e.g., vector, square, cube, etc.)
+    footprint : ndarray, tuple of int, or None
+        If ``None``, ``footprint`` will be a N-D array with 3 elements for each
+        dimension (e.g., vector, square, cube, etc.). If `footprint` is a
+        tuple of integers, it will be an array of ones with the given shape.
+        Otherwise, if ``behavior=='rank'``, ``footprint`` is a 2-D array of 1's
+        and 0's. If ``behavior=='ndimage'``, ``footprint`` is a N-D array of
+        1's and 0's with the same number of dimension as ``image``.
+        Note that upstream scikit-image currently does not support supplying
+        a tuple for `footprint`. It is added here to avoid overhead of
+        generating a small weights array in cases where it is not needed.
     out : ndarray, (same dtype as image), optional
         If None, a new array is allocated.
     mode : {'reflect', 'constant', 'nearest', 'mirror','‘wrap'}, optional
@@ -121,12 +125,21 @@ def median(image, footprint=None, out=None, mode='nearest', cval=0.0,
         # return generic.median(image, selem=selem, out=out)
 
     if footprint is None:
-        footprint = ndi.generate_binary_structure(image.ndim, image.ndim)
+        footprint_shape = (3,) * image.ndim
+    elif isinstance(footprint, tuple):
+        if len(footprint) != image.ndim:
+            raise ValueError("tuple footprint must have ndim matching image")
+        footprint_shape = footprint
+        footprint = None
+    else:
+        footprint_shape = footprint.shape
 
     if algorithm == 'sorting':
         can_use_histogram = False
     elif algorithm in ['auto', 'histogram']:
-        can_use_histogram, reason = _can_use_histogram(image, footprint)
+        can_use_histogram, reason = _can_use_histogram(
+            image, footprint, footprint_shape
+        )
     else:
         raise ValueError(f"unknown algorithm: {algorithm}")
 
@@ -142,7 +155,7 @@ def median(image, footprint=None, out=None, mode='nearest', cval=0.0,
     use_histogram = can_use_histogram
     if algorithm == 'auto':
         # prefer sorting-based algorithm if footprint shape is small
-        use_histogram = use_histogram and prod(footprint.shape) > 150
+        use_histogram = use_histogram and prod(footprint_shape) > 150
 
     if use_histogram:
         try:
@@ -162,8 +175,16 @@ def median(image, footprint=None, out=None, mode='nearest', cval=0.0,
 
             # TODO: Can't currently pass an output array into _median_hist as a
             #       new array currently needs to be created during padding.
-            temp = _median_hist(image, footprint, mode=mode, cval=cval,
-                                **algorithm_kwargs)
+
+            # pass shape if explicit footprint isn't needed
+            # (use new variable name in case KernelResourceError occurs)
+            temp = _median_hist(
+                image,
+                footprint_shape if footprint is None else footprint,
+                mode=mode,
+                cval=cval,
+                **algorithm_kwargs
+            )
             if output_array_provided:
                 out[:] = temp
             else:
@@ -175,12 +196,13 @@ def median(image, footprint=None, out=None, mode='nearest', cval=0.0,
             # Fall back to sorting-based implementation if we encounter a
             # resource limit (e.g. insufficient shared memory per block).
             warn("Kernel resource error encountered in histogram-based "
-                 f"median kerne: {e}\n"
+                 f"median kernel: {e}\n"
                  "Falling back to sorting-based median instead.")
 
     if algorithm_kwargs:
         warn(f"algorithm_kwargs={algorithm_kwargs} ignored for sorting-based "
              f"algorithm")
 
-    return ndi.median_filter(image, footprint=footprint, output=out, mode=mode,
-                             cval=cval)
+    size = footprint_shape if footprint is None else None
+    return ndi.median_filter(image, size=size, footprint=footprint, output=out,
+                             mode=mode, cval=cval)

python/cucim/src/cucim/skimage/filters/_median_hist.py

@@ -230,7 +230,7 @@ def _check_global_scratch_space_size(
     return (n_fine + n_coarse) * cp.dtype(hist_dtype).itemsize
 
 
-def _can_use_histogram(image, footprint):
+def _can_use_histogram(image, footprint, footprint_shape=None):
     """Validate compatibility with histogram-based median.
 
     Parameters
@@ -255,23 +255,32 @@ def _can_use_histogram(image, footprint):
         return False, "Only 8 and 16-bit integer image types (signed or "
         "unsigned)."
 
+    if footprint is None:
+        if footprint_shape is None:
+            raise ValueError(
+                "must provide either footprint or footprint_shape"
+            )
+    else:
+        footprint_shape = footprint.shape
+
     # only odd-sized footprints are supported
-    if not all(s % 2 == 1 for s in footprint.shape):
+    if not all(s % 2 == 1 for s in footprint_shape):
         return False, "footprint must have odd size on both axes"
 
-    if any(s == 1 for s in footprint.shape):
+    if any(s == 1 for s in footprint_shape):
         return False, "footprint must have size >= 3"
 
     # footprint radius can't be larger than the image
     # TODO: need to check if we need this exact restriction
     #       (may be specific to OpenCV's boundary handling)
-    radii = tuple(s // 2 for s in footprint.shape)
+    radii = tuple(s // 2 for s in footprint_shape)
     if any(r > s for r, s in zip(radii, image.shape)):
         return False, "footprint half-width cannot exceed the image extent"
 
-    # only fully populated footprint is supported
-    if not np.all(footprint):  # synchronizes!
-        return False, "footprint must be 1 everywhere"
+    if footprint is not None:
+        # only fully populated footprint is supported
+        if not np.all(footprint):  # synchronizes!
+            return False, "footprint must be 1 everywhere"
 
     return True, None
 
@@ -439,7 +448,19 @@ def _median_hist(image, footprint, output=None, mode='mirror', cval=0,
             "Use of a user-defined output array has not been implemented"
         )
 
-    compatible_image, reason = _can_use_histogram(image, footprint)
+    if footprint is None:
+        footprint_shape = (3,) * image.ndim
+        med_pos = prod(footprint_shape) // 2
+    elif isinstance(footprint, tuple):
+        footprint_shape = footprint
+        footprint = None
+        med_pos = prod(footprint_shape) // 2
+    else:
+        footprint_shape = footprint.shape
+        med_pos = footprint.size // 2
+    compatible_image, reason = _can_use_histogram(
+        image, footprint, footprint_shape
+    )
     if not compatible_image:
         raise ValueError(reason)
 
@@ -451,13 +472,12 @@ def _median_hist(image, footprint, output=None, mode='mirror', cval=0,
     if image.dtype.kind not in 'iu':
         raise ValueError("only integer-type images are accepted")
 
-    radii = tuple(s // 2 for s in footprint.shape)
+    radii = tuple(s // 2 for s in footprint_shape)
     # med_pos is the index corresponding to the median
     # (calculation here assumes all elements of the footprint are True)
-    med_pos = footprint.size // 2
 
     params = _get_kernel_params(
-        image, footprint.shape, value_range, partitions
+        image, footprint_shape, value_range, partitions
     )
 
     # pad as necessary to avoid boundary artifacts

python/cucim/src/cucim/skimage/filters/tests/test_median.py

@@ -69,16 +69,22 @@ def test_selem_kwarg_deprecation(image):
         (3, 3), (5, 5), (9, 15), (2, 2), (1, 1), (2, 7), (23, 23), (15, 35),
     ]
 )
+@pytest.mark.parametrize('footprint_tuple', (False, True))
 @pytest.mark.parametrize('out', [None, cp.uint8, cp.float32, 'array'])
-def test_median_behavior(camera, behavior, func, mode, footprint_shape, out):
-    footprint = cp.ones(footprint_shape, dtype=bool)
+def test_median_behavior(
+    camera, behavior, func, mode, footprint_shape, footprint_tuple, out
+):
+    if footprint_tuple:
+        footprint = footprint_shape
+    else:
+        footprint = cp.ones(footprint_shape, dtype=bool)
     cam2 = camera[:, :177]  # use anisotropic size
     assert cam2.dtype == cp.uint8
     if out == 'array':
         out = cp.zeros_like(cam2)
     assert_allclose(
         median(cam2, footprint, mode=mode, behavior=behavior, out=out),
-        func(cam2, size=footprint.shape, mode=mode, output=out),
+        func(cam2, size=footprint_shape, mode=mode, output=out),
     )