Python NumPy花式索引-从不同通道裁剪不同的ROI

假设我们有以下numpy 4D阵列（3x1x4x4）：

现在，我想在不同的位置以相同的大小裁剪n个子阵列中的每一个子阵列：

size = 2
locations = np.array([
    [0, 1],
    [1, 1],
    [0, 2]
])

执行此操作的慢方法如下所示：

crops = np.stack([d[:, y:y+size, x:x+size] 
     for d, (y,x) in zip(data, locations)])

>>> array([[[[ 1,  2],
             [ 5,  6]]],

           [[[21, 22],
             [25, 26]]],

           [[[34, 35],
             [38, 39]]]])

现在，我正在寻找一种方法来实现这一点与numpy的花式索引。我已经花了几个小时想办法解决这个问题。我是否忽略了解决这个问题的简单方法？有没有一些numpy索引专家可以帮助我？

我们可以扩展到您的

3D

案例，利用基于

的滑动窗口视图进行高效的补丁提取，就像这样-

from skimage.util.shape import view_as_windows

def get_patches(data, locations, size):
    # Get 2D sliding windows for each element off data
    w = view_as_windows(data, (1,1,size,size))
    
    # Use fancy/advanced indexing to select the required ones
    return w[np.arange(len(locations)), :, locations[:,0], locations[:,1]][:,:,0,0]

我们需要那些
1,1
作为窗口参数来
查看窗口，因为它期望窗口的元素数量与输入数据的DIM数量相同。我们沿着
数据的最后两个轴滑动，因此将前两个轴保持为
1s
，基本上不沿着
数据的前两个轴滑动

一个通道和多个通道数据的示例运行-

In [78]: n, c, h, w = 3, 1, 4, 4 # number of channels = 1
    ...: data = np.arange(n * c * h * w).reshape(n, c, h, w)
    ...: 
    ...: size = 2
    ...: locations = np.array([
    ...:     [0, 1],
    ...:     [1, 1],
    ...:     [0, 2]
    ...: ])
    ...: 
    ...: crops = np.stack([d[:, y:y+size, x:x+size] 
    ...:      for d, (y,x) in zip(data, locations)])

In [79]: print np.allclose(get_patches(data, locations, size), crops)
True

In [80]: n, c, h, w = 3, 5, 4, 4 # number of channels = 5
    ...: data = np.arange(n * c * h * w).reshape(n, c, h, w)
    ...: 
    ...: size = 2
    ...: locations = np.array([
    ...:     [0, 1],
    ...:     [1, 1],
    ...:     [0, 2]
    ...: ])
    ...: 
    ...: crops = np.stack([d[:, y:y+size, x:x+size] 
    ...:      for d, (y,x) in zip(data, locations)])

In [81]: print np.allclose(get_patches(data, locations, size), crops)
True


标杆管理
其他办法-

# Original soln
def stack(data, locations, size):
    crops = np.stack([d[:, y:y+size, x:x+size] 
         for d, (y,x) in zip(data, locations)])    
    return crops

# scholi's soln
def allocate_assign(data, locations, size):
    n, c, h, w = data.shape
    crops = np.zeros((n,c,size,size))
    for i, (y,x) in enumerate(locations):
        crops[i,0,:,:] = data[i,0,y:y+size,x:x+size]
    return crops

从评论来看，OP似乎对数据为
（512,1,60,60）
且
大小为
12,24,48
的情况感兴趣。那么，让我们用一个函数来设置数据-

# Setup data
def create_inputs(size):
    np.random.seed(0)
    n, c, h, w = 512, 1, 60, 60
    data = np.arange(n * c * h * w).reshape(n, c, h, w)
    locations = np.random.randint(0,3,(n,2))
    return data, locations, size

时间安排-

In [186]: data, locations, size = create_inputs(size=12)

In [187]: %timeit stack(data, locations, size)
     ...: %timeit allocate_assign(data, locations, size)
     ...: %timeit get_patches(data, locations, size)
1000 loops, best of 3: 1.26 ms per loop
1000 loops, best of 3: 1.06 ms per loop
10000 loops, best of 3: 124 µs per loop

In [188]: data, locations, size = create_inputs(size=24)

In [189]: %timeit stack(data, locations, size)
     ...: %timeit allocate_assign(data, locations, size)
     ...: %timeit get_patches(data, locations, size)
1000 loops, best of 3: 1.66 ms per loop
1000 loops, best of 3: 1.55 ms per loop
1000 loops, best of 3: 470 µs per loop

In [190]: data, locations, size = create_inputs(size=48)

In [191]: %timeit stack(data, locations, size)
     ...: %timeit allocate_assign(data, locations, size)
     ...: %timeit get_patches(data, locations, size)
100 loops, best of 3: 2.8 ms per loop
100 loops, best of 3: 3.33 ms per loop
1000 loops, best of 3: 1.45 ms per loop

堆叠速度很慢。由于大小已知，最好先分配裁剪后的数组

crops = np.zeros((3,1,size,size))
for i, (y,x) in enumerate(locations):
    crops[i,0,:,:] = data[i,0,y:y+size,x:x+size]

Divakar的解决方案在速度上是最差的。我使用%%timeit获得92.3µs
您的堆栈解决方案是35.4μs，我提供的示例得到了29.3μs。
因为这是一个有趣的问题，所以进行了投票。看看这是否能解决问题-。好的，我来看看帖子谢谢你的快速回复！我还在读关于视窗的文件，并在尝试，正如我看到你的答案。我会试试看，再次感谢！：）有趣的事实：一切都很好，但是：两种解决方案都同样快速。。。此外，我在GPU上与cupy一起工作，我不得不将这些裁剪操作切换到CPU，因为view_as_windows只是CPU。因此，两种实现的执行速度都提高了3倍。只是想与您分享这两个解决方案，就像基于堆栈的解决方案和此解决方案一样？我认为，要看到此解决方案的性能提升，您需要更多的窗口，即
位置中有大量的行。是的，两种解决方案在CPU上的速度都比在GPU上快（尽管我无法在GPU上测试您的解决方案）。我现在的n是20，我可以试试更大的，比如512左右…@BloodyD后端到底是什么？不要认为NumPy本机支持GPU。我怀疑，普通的python循环比列表理解更快。尝试将第一个维度增加到维度？或者np.stack速度太慢了？好的，我已经测试了你的解决方案，第一维度设置为512，与我的np.stack解决方案相比，每次分配数组都会有一个小的提升。然而，Divakar的解决方案稍微快一点。
crops = np.zeros((3,1,size,size))
for i, (y,x) in enumerate(locations):
    crops[i,0,:,:] = data[i,0,y:y+size,x:x+size]