在 3D 数组的每个 2D 切片上的随机位置填充块

Question

我有 3D numpy 数组，例如，像这样：

array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[16, 17, 18, 19],
        [20, 21, 22, 23],
        [24, 25, 26, 27],
        [28, 29, 30, 31]]])

有没有一种方法可以索引它，例如，我选择第一个平面中 2x2 元素的右上角，以及第二个平面中的中心 2x2 元素子数组？这样我就可以将元素 2,3,6,7,21,22,25,26 归零：

array([[[ 0,  1,  0,  0],
        [ 4,  5,  0,  0],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[16, 17, 18, 19],
        [20,  0,  0, 23],
        [24,  0,  0, 27],
        [28, 29, 30, 31]]])

我有一批图像，我需要将一个固定大小的小窗口归零，但在批次中的每个图像的不同（随机）位置。第一个维度是图像的数量。

类似这样的： a[:, x: x+2, y: y+2] = 0

其中 x 和 y 是向量，对于 a 的每个第一维都有不同的值。

Answer 1

方法#1：这是一种主要基于linear-indexing的方法-

def random_block_fill_lidx(a, N, fillval=0):
    # a is input array
    # N is blocksize

    # Store shape info
    m,n,r = a.shape

    # Get all possible starting linear indices for each 2D slice
    possible_start_lidx = (np.arange(n-N+1)[:,None]*r + range(r-N+1)).ravel()

    # Get random start indices from all possible ones for all 2D slices
    start_lidx = np.random.choice(possible_start_lidx, m)

    # Get linear indices for the block of (N,N)
    offset_arr = (a.shape[-1]*np.arange(N)[:,None] + range(N)).ravel()

    # Add in those random start indices with the offset array
    idx = start_lidx[:,None] + offset_arr

    # On a 2D view of the input array, use advance-indexing to set fillval.
    a.reshape(m,-1)[np.arange(m)[:,None], idx] = fillval
    return a

方法 #2： 这是另一种可能更有效的方法（用于大型 2D 切片），使用 advanced-indexing -

def random_block_fill_adv(a, N, fillval=0):
    # a is input array
    # N is blocksize

    # Store shape info
    m,n,r = a.shape

    # Generate random start indices for second and third axes keeping proper
    # distance from the boundaries for the block to be accomodated within.
    idx0 = np.random.randint(0,n-N+1,m)
    idx1 = np.random.randint(0,r-N+1,m)

    # Setup indices for advanced-indexing.

    # First axis indices would be simply the range array to select one per elem.
    # We need to extend this to 3D so that the latter dim indices could be aligned.
    dim0 = np.arange(m)[:,None,None]

    # Second axis indices would idx0 with broadcasted additon of blocksized 
    # range array to cover all block indices along this axis. Repeat for third.
    dim1 = idx0[:,None,None] + np.arange(N)[:,None]
    dim2 = idx1[:,None,None] + range(N)
    a[dim0, dim1, dim2] = fillval
    return a

方法 #3： 使用旧的可信赖循环 -

def random_block_fill_loopy(a, N, fillval=0):
    # a is input array
    # N is blocksize

    # Store shape info
    m,n,r = a.shape

    # Generate random start indices for second and third axes keeping proper
    # distance from the boundaries for the block to be accomodated within.
    idx0 = np.random.randint(0,n-N+1,m)
    idx1 = np.random.randint(0,r-N+1,m)

    # Iterate through first and use slicing to assign fillval.
    for i in range(m):
        a[i, idx0[i]:idx0[i]+N, idx1[i]:idx1[i]+N] = fillval        
    return a

示例运行 -

In [357]: a = np.arange(2*4*7).reshape(2,4,7)

In [358]: a
Out[358]: 
array([[[ 0,  1,  2,  3,  4,  5,  6],
        [ 7,  8,  9, 10, 11, 12, 13],
        [14, 15, 16, 17, 18, 19, 20],
        [21, 22, 23, 24, 25, 26, 27]],

       [[28, 29, 30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39, 40, 41],
        [42, 43, 44, 45, 46, 47, 48],
        [49, 50, 51, 52, 53, 54, 55]]])

In [359]: random_block_fill_adv(a, N=3, fillval=0)
Out[359]: 
array([[[ 0,  0,  0,  0,  4,  5,  6],
        [ 7,  0,  0,  0, 11, 12, 13],
        [14,  0,  0,  0, 18, 19, 20],
        [21, 22, 23, 24, 25, 26, 27]],

       [[28, 29, 30, 31, 32, 33, 34],
        [35, 36, 37, 38,  0,  0,  0],
        [42, 43, 44, 45,  0,  0,  0],
        [49, 50, 51, 52,  0,  0,  0]]])

有趣的东西： 就地填充，如果我们继续运行random_block_fill_adv(a, N=3, fillval=0)，我们最终会以全零结束a。因此，还要验证代码。

---

运行时测试

In [579]: a = np.random.randint(0,9,(10000,4,4))

In [580]: %timeit random_block_fill_lidx(a, N=2, fillval=0)
     ...: %timeit random_block_fill_adv(a, N=2, fillval=0)
     ...: %timeit random_block_fill_loopy(a, N=2, fillval=0)
     ...: 
1000 loops, best of 3: 545 µs per loop
1000 loops, best of 3: 891 µs per loop
100 loops, best of 3: 10.6 ms per loop

In [581]: a = np.random.randint(0,9,(1000,40,40))

In [582]: %timeit random_block_fill_lidx(a, N=10, fillval=0)
     ...: %timeit random_block_fill_adv(a, N=10, fillval=0)
     ...: %timeit random_block_fill_loopy(a, N=10, fillval=0)
     ...: 
1000 loops, best of 3: 739 µs per loop
1000 loops, best of 3: 671 µs per loop
1000 loops, best of 3: 1.27 ms per loop

所以，选择哪一个取决于第一个轴的长度和块大小。