从数组中切出一块并用零填充

Question

我需要从 N-dim 数组中剪切给定大小和给定位置的一部分。 如果零件太大，我需要用零填充它以达到给定的大小。

为简单起见，这些示例都是 2D 的。

给定的矩阵：

[[1 8 3 3 8]
 [5 8 6 7 6]
 [8 3 5 6 5]
 [2 6 2 4 6]
 [6 5 3 7 4]]

我想切割 [2,4] 部分，从索引 (1,2) 开始， 我剪的部分不够大，所以填充 需要零。 想要的结果：

[[6 7 6 0]
 [5 6 5 0]]

我设法编写丑陋而不是 N-dim 的代码来做到这一点。

# set example numbers
matrix =  numpy.random.randint(low=1, high=9, size=(5,5))
matrix_size = np.array(matrix.shape)

# size of the part we want to have in the end
size = np.array([2, 4])
# starting point of the cut
mini = [1, 2]

#calculating max index (in the given matrix) for the part we want to cut
maxi = np.add(size - 1 , mini)
cut_max_ind = np.minimum(maxi, matrix_size - 1) + 1

# copy from matrix to cut
# ??? a way to generalize it for N-dim ???
cut = matrix[mini[0]:cut_max_ind[0], mini[1]:cut_max_ind[1]]

#culculate the padding size
padding =  np.add(matrix_size - 1, maxi*-1)
padding_size = np.minimum(np.zeros((matrix.ndim), dtype=np.uint8), padding) * -1

for j in range(0, matrix.ndim):

    if (padding_size[j]):
        pad_width = size
        pad_width[j] = padding_size[j]
        pad_pice = np.zeros((pad_width), dtype = np.uint8)
        cut = np.append(cut, pad_pice, axis = j)

print "matrix"
print matrix
print "cut"
print cut

任何改进和概括的想法？

Answer 1

您可以通过预先分配零数组然后修改切片以满足您的需要来更轻松地解决它：

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

def extract_piece(array, in_idx):
    # make sure number of dimensions match
    assert array.ndim == len(in_idx)
    # preallocate output array
    out = numpy.zeros([i.stop - i.start for i in in_idx], dtype=array.dtype)

    # modify reading slices to not exceed bounds
    in_idx = [slice(i.start, min(i.stop, s), i.step) for s, i in zip(array.shape, in_idx)]
    # modify writing slices to fit size of read data
    out_idx = [slice(0, i.stop - i.start) for i in in_idx]

    # Copy data
    out[out_idx] = array[in_idx]
    return out

print a
print extract_piece(a, (slice(1, 3), slice(2, 6)))

或者在 4D 示例中

extract_piece(
    numpy.random.rand(4, 4, 4, 4),  # 4D data
    (
        slice(0,2),  # 1st dimension
        slice(0,4),  # 2nd dimension
        slice(0,6),  # 3rd dimension
        slice(0,1),  # 4th dimension
    )
)

如果您不熟悉切片：

a[1:2]

和

一样

a[slice(1,2)]