python - numpy 数组 : efficient way to compute argmax within a fixed window at a set of rows and columns given as input

Question

假设我们有一个二维数组 arr，一组由 rows 和 cols 定义的位置，以及一个形状为 (5, 1)。对于(i, j)，我们需要arr[i-2:i+2, j]中最大值的索引。我们想对所有输入 (i, j) 对重复。

import numpy as np

def get_max_index(arr, i, j, stride):
    nr, nc = arr.shape
    i_low = max(0, i - stride)
    i_up = min(i + stride, nr)

    idx = np.argmax(arr[i_low:i_up, j])

    # transform back to original index.
    return i_low + idx, j

# Given numpy array
arr = np.array([
    [1,2,3,6,7],
    [4,5,6,15,8],
    [7,8,9,24,9],
    [1,1,1,3,10],
    [2,2,2,6,11],
    [3,3,3,9,12],
    [4,4,4,4,42]
    ])

# Rows and columns at which the windows will be centered.
rows = np.array([2, 4, 6, 6])
cols = np.array([1, 1, 3, 4])

# Measure corresponding to window of size 5
stride = 2

# Apply the function on the input rows and cols.
res = [get_max_index(arr, i, j, stride) for i, j in zip(rows, cols)]

assert res == [(2, 1), (2, 1), (5, 3), (6, 4)]

我很好奇是否有更快的 numpy 方法来代替使用列表理解。

它与“morphological dilation”有些相似，但它位于数组单元格的子集上，我们需要索引。

Answer 1

我们可以利用 np.lib.stride_tricks.as_strided基于 scikit-image's view_as_windows将滑动窗口 View 放入输入的最小值填充版本(以考虑边界情况)，然后在其中获取 argmax 值并针对给定行进行偏移。

因此，实现将是 -

from skimage.util.shape import view_as_windows

def windowed_argmax(arr, rows, cols, stride):
    # Get full window extent 
    W = 2*stride+1

    #  Pad with minimum value, so that on boundaries we will skip those
    a = np.pad(arr,((stride,stride),(0,0)),'constant',constant_values=arr.min()-1)

    # Get sliding windows
    w = view_as_windows(a,(W,1))[...,0]

    # Index into those specific rows, cols positions; get argmax, offset back
    return np.c_[rows+w[rows,cols].argmax(1)-stride,cols]

sample 运行-

In [75]: arr
Out[75]: 
array([[ 1,  2,  3,  6,  7],
       [ 4,  5,  6, 15,  8],
       [ 7,  8,  9, 24,  9],
       [ 1,  1,  1,  3, 10],
       [ 2,  2,  2,  6, 11],
       [ 3,  3,  3,  9, 12],
       [ 4,  4,  4,  4, 42]])

In [76]: rows
Out[76]: array([2, 4, 6, 6])

In [77]: cols
Out[77]: array([1, 1, 3, 4])

In [78]: windowed_argmax(arr, rows, cols, stride=2)
Out[78]: 
array([[2, 1],
       [2, 1],
       [5, 3],
       [6, 4]])