通过严格比较对函数进行向量化，以在 2D 数组中查找局部最小值和最大值

我正在尝试提高返回输入 2D NumPy 数组的局部最小值和最大值的函数的性能。该函数按预期工作，但对于我的用例来说太慢了。我想知道是否可以创建此函数的矢量化版本以提高其性能。

Here is the formal definition for defining whether an element is a local minima (maxima):

where

A=[a_m,n]是二维矩阵，m and n分别是行和列，w_h and w_w分别是滑动窗口的高度和宽度。

我尝试过使用skimage.morphology.local_minimum and skimage.morphology.local_maxima，但如果某个元素的值小于或等于（大于或等于）其所有邻居，则他们将其视为最小值（最大值）。
就我而言，如果一个元素严格小于（大于）其所有邻居，我需要该函数将其视为最小（最大）。

当前的实现使用滑动窗口方法numpy.lib.stride_tricks.sliding_window_view，但函数不一定非要使用这种方式。

这是我当前的实现：

import numpy as np

def get_local_extrema(array, window_size=(3, 3)):
    # Check if the window size is valid
    if not all(size % 2 == 1 and size >= 3 for size in window_size):
        raise ValueError("Window size must be odd and >= 3 in both dimensions.")

    # Create a map to store the local minima and maxima
    minima_map = np.zeros_like(array)
    maxima_map = np.zeros_like(array)

    # Save the shape and dtype of the original array for later
    original_size = array.shape
    original_dtype = array.dtype
    # Get the halved window size
    half_window_size = tuple(size // 2 for size in window_size)

    # Pad the array with NaN values to handle the edge cases
    padded_array = np.pad(array.astype(float),
                         tuple((size, size) for size in half_window_size),
                         mode='constant', constant_values=np.nan)

    # Generate all the sliding windows
    windows = np.lib.stride_tricks.sliding_window_view(padded_array, window_size).reshape(
        original_size[0] * original_size[1], *window_size)

    # Create a mask to ignore the central element of the window
    mask = np.ones(window_size, dtype=bool)
    mask[half_window_size] = False

    # Iterate through all the windows
    for i in range(windows.shape[0]):
        window = windows[i]
        # Get the value of the central element
        center_val = window[half_window_size]
        # Apply the mask to ignore the central element
        masked_window = window[mask]

        # Get the row and column indices of the central element
        row = i // original_size[1]
        col = i % original_size[1]

        # Check if the central element is a local minimum or maximum
        if center_val > np.nanmax(masked_window):
            maxima_map[row, col] = center_val
        elif center_val < np.nanmin(masked_window):
            minima_map[row, col] = center_val

    return minima_map.astype(original_dtype), maxima_map.astype(original_dtype)

a = np.array([[8, 8, 4, 1, 5, 2, 6, 3],
              [6, 3, 2, 3, 7, 3, 9, 3],
              [7, 8, 3, 2, 1, 4, 3, 7],
              [4, 1, 2, 4, 3, 5, 7, 8],
              [6, 4, 2, 1, 2, 5, 3, 4],
              [1, 3, 7, 9, 9, 8, 7, 8],
              [9, 2, 6, 7, 6, 8, 7, 7],
              [8, 2, 1, 9, 7, 9, 1, 1]])

(minima, maxima) = get_local_extrema(a)

print(minima)
# [[0 0 0 1 0 2 0 0]
#  [0 0 0 0 0 0 0 0]
#  [0 0 0 0 1 0 0 0]
#  [0 1 0 0 0 0 0 0]
#  [0 0 0 1 0 0 3 0]
#  [1 0 0 0 0 0 0 0]
#  [0 0 0 0 6 0 0 0]
#  [0 0 1 0 0 0 0 0]]

print(maxima)
# [[0 0 0 0 0 0 0 0]
#  [0 0 0 0 7 0 9 0]
#  [0 8 0 0 0 0 0 0]
#  [0 0 0 4 0 0 0 8]
#  [6 0 0 0 0 0 0 0]
#  [0 0 0 0 0 0 0 8]
#  [9 0 0 0 0 0 0 0]
#  [0 0 0 9 0 9 0 0]]

expected_minima = np.array([[0, 0, 0, 1, 0, 2, 0, 0],
                            [0, 0, 0, 0, 0, 0, 0, 0],
                            [0, 0, 0, 0, 1, 0, 0, 0],
                            [0, 1, 0, 0, 0, 0, 0, 0],
                            [0, 0, 0, 1, 0, 0, 3, 0],
                            [1, 0, 0, 0, 0, 0, 0, 0],
                            [0, 0, 0, 0, 6, 0, 0, 0],
                            [0, 0, 1, 0, 0, 0, 0, 0]])

expected_maxima = np.array([[0, 0, 0, 0, 0, 0, 0, 0],
                            [0, 0, 0, 0, 7, 0, 9, 0],
                            [0, 8, 0, 0, 0, 0, 0, 0],
                            [0, 0, 0, 4, 0, 0, 0, 8],
                            [6, 0, 0, 0, 0, 0, 0, 0],
                            [0, 0, 0, 0, 0, 0, 0, 8],
                            [9, 0, 0, 0, 0, 0, 0, 0],
                            [0, 0, 0, 9, 0, 9, 0, 0]])

np.testing.assert_array_equal(minima, expected_minima)
np.testing.assert_array_equal(maxima, expected_maxima)

print('All tests passed')

任何有关如何向量化此函数的建议或想法将不胜感激。

提前致谢！

EDIT #1
在玩了一下 NumPy 后，如果我理解正确的话，我设法让以下代码几乎以完全矢量化的方式工作：

def get_local_extrema_2(img):
  minima_map = np.zeros_like(img)
  maxima_map = np.zeros_like(img)

  minima_map[1:-1, 1:-1] = np.where(
    (a[1:-1, 1:-1] < a[:-2, 1:-1]) &
    (a[1:-1, 1:-1] < a[2:, 1:-1]) &
    (a[1:-1, 1:-1] < a[1:-1, :-2]) &
    (a[1:-1, 1:-1] < a[1:-1, 2:]) &
    (a[1:-1, 1:-1] < a[2:, 2:]) &
    (a[1:-1, 1:-1] < a[:-2, :-2]) &
    (a[1:-1, 1:-1] < a[2:, :-2]) &
    (a[1:-1, 1:-1] < a[:-2, 2:]),
    a[1:-1, 1:-1],
    0)
  
  maxima_map[1:-1, 1:-1] = np.where(
    (a[1:-1, 1:-1] > a[:-2, 1:-1]) &
    (a[1:-1, 1:-1] > a[2:, 1:-1]) &
    (a[1:-1, 1:-1] > a[1:-1, :-2]) &
    (a[1:-1, 1:-1] > a[1:-1, 2:]) &
    (a[1:-1, 1:-1] > a[2:, 2:]) &
    (a[1:-1, 1:-1] > a[:-2, :-2]) &
    (a[1:-1, 1:-1] > a[2:, :-2]) &
    (a[1:-1, 1:-1] > a[:-2, 2:]),
    a[1:-1, 1:-1],
    0)

  return minima_map, maxima_map

get_local_extrema_2 的输出是：
最小地图：

[[0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0]
 [0 0 0 0 1 0 0 0]
 [0 1 0 0 0 0 0 0]
 [0 0 0 1 0 0 3 0]
 [0 0 0 0 0 0 0 0]
 [0 0 0 0 6 0 0 0]
 [0 0 0 0 0 0 0 0]]

千里马地图：

[[0 0 0 0 0 0 0 0]
 [0 0 0 0 7 0 9 0]
 [0 8 0 0 0 0 0 0]
 [0 0 0 4 0 0 0 0]
 [0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0]]

上述问题是未检测到边界上的最小值或最大值的像素。

EDIT #2
即使输出数组中有 1 而不是局部最小值（最大值）的值，即 0 和 1（或 False 和 True）的二维数组，也没关系。

EDIT #3
这是该函数的一个版本基于克里斯·卢恩戈 https://stackoverflow.com/users/7328782/cris-luengo's answer https://stackoverflow.com/a/75049306/5495385。请注意使用“镜像”模式（相当于 NumPy 的“反射”），​​这样如果最小值或最大值位于边缘，它就不会被复制到边界之外，并且会脱颖而出。这样，就不需要用矩阵的最小或最大元素填充图像。我认为这是完成此任务的最有效的方法：

import numpy as np
import scipy

def get_local_extrema_v3(image):
    footprint = np.ones((3, 3), dtype=bool)
    footprint[1, 1] = False
    minima = image * (scipy.ndimage.grey_erosion(image, footprint=footprint, mode='mirror') > image)
    maxima = image * (scipy.ndimage.grey_dilation(image, footprint=footprint, mode='mirror') < image)
    return minima, maxima

您对局部最大值的定义是有缺陷的。例如，在一维数组中[1,2,3,4,4,3,2,1]，存在局部最大值，但您的定义忽略了它。skimage.morphology.local_maxima将正确识别该局部最大值。

如果您确实需要实现您的定义，我将使用带有窗口大小的方形结构元素的膨胀（侵蚀），但不包括中心像素。原始图像中比滤波图像中更大（更小）的任何像素都将满足局部最大值（最小值）的定义。

我使用 scikit-image 实现了这一点，但发现它在图像边缘做了奇怪的事情，因此它不会检测边缘附近的局部最大值或最小值：

se = np.ones((3, 3))
se[1, 1] = 0
minima = a * (skimage.morphology.erosion(a, footprint=se) > a)
maxima = a * (skimage.morphology.dilation(a, footprint=se) < a)

使用 DIPlib （披露：我是作者）这也可以在图像边缘正常工作：

import diplib as dip

se = np.ones((3, 3), dtype=np.bool_)
se[1, 1] = False
minima = a * (dip.Erosion(a, se) > a)
maxima = a * (dip.Dilation(a, se) < a)

^{Looking at the source code for skimage.morphology.dilation, it calls scipy.ndimage.grey_dilation with the default boundary extension, which is 'reflect'. This means that every local maximum at the image edge will have a neighbor with the same value, and hence not detected as local maximum in this definition. Instead, it should use the 'constant' extension, with cval set to the minimum possible value for the data type. For example, for an uint8 input array, it should do ndi.grey_dilation(image, footprint=footprint, output=out, mode='constant', cval=0). GitHub issue https://github.com/scikit-image/scikit-image/issues/6665}