1

我有两个输入数组：维度 (i,j,k) 的 data_arr 和维度 (i,j) 的 index_arr。index_arr 中的条目是 [0, k-1] 范围内的整数。我想创建一个维度为 (i,j) 的输出数组 (output_arr)，其中对于 output_arr 的每个元素，index_arr 告诉我可以选择哪些元素。

换句话说 output_arr[i,j] = data_arr[i,j, index_arr[i, j]]

显然，我可以用双 for 循环以冰冷的速度做到这一点。我更喜欢使用智能索引的更快捷的东西。目前我能设计的最好的方法是创建两个额外的大小为 (i,j) 的二维矩阵。

下面是一个简单的 MWE，它使用标准拜耳模式从 RGB 图像创建马赛克图像。我希望能够摆脱X_ind和Y_ind

import numpy as np
import time


if __name__ == '__main__':
    img_width = 1920
    img_height = 1080
    img_num_colours = 3

    red_arr = np.ones([img_height, img_width], dtype=np.uint16) * 10
    green_arr = np.ones([img_height, img_width], dtype=np.uint16) * 20
    blue_arr = np.ones([img_height, img_width], dtype=np.uint16) * 30

    img_arr = np.dstack((red_arr, green_arr, blue_arr))

    bayer_arr = np.ones([img_height, img_width], dtype=np.uint16)
    bayer_arr[0::2,0::2] = 0 # Red entries in bater patter
                             # Green entries are already set by np.ones intialisation
    bayer_arr[1::2,1::2] = 2 # blue entries in bayer patter
    print("bayer\n",bayer_arr[:8,:12], "\n")

    mosaiced_arr = np.zeros([img_height, img_width], dtype=np.uint16)
    Y_ind = np.repeat(np.arange(0, img_width).reshape(1, img_width), img_height, 0)
    X_ind = np.repeat(np.arange(0, img_height).reshape(img_height, 1), img_width, 1)

    start_time = time.time()
    demos_arr = img_arr[X_ind, Y_ind, bayer_arr]
    end_time = time.time()

    print(demos_arr.shape)
    print("demos\n",demos_arr[:8,:12], "\n")
    print("Mosaic took {:.3f}s".format(end_time - start_time))

编辑：正如@Georgy 所指出的，这个问题类似于我在搜索词中没有找到的这个问题，所以也许这篇文章将作为那个帖子的标志。另一篇文章中的答案是适用的，尽管扁平索引算法不同，因为我的维度排序不同。上面的答案等同ogrid于另一个问题中的版本。实际上ogrid可以通过替换对代码进行以下更改来使用：

# Y_ind = np.repeat(np.arange(0, img_width).reshape(1, img_width), img_height, 0)
# X_ind = np.repeat(np.arange(0, img_height).reshape(img_height, 1), img_width, 1)
X_ind, Y_ind = np.ogrid[0:img_height, 0:img_width]

您可以像这样实现选择选项（仅限于在 32 个选项之间进行选择）：

start_time = time.time()
demos_arr = bayer_arr.choose((img_arr[...,0], img_arr[...,1], img_arr[...,2]))
end_time = time.time()

在我的机器上，ogrid 解决方案在 12 毫秒内运行，而选择解决方案在 34 毫秒内运行

4

1 回答 1

2

你想要numpy.take_along_axis：

output_arr = numpy.take_along_axis(data_arr, index_arr[:, :, numpy.newaxis], axis=2)
output_arr = output_arr[:,:,0]  # Since take_along_axis keeps the same number of dimensions

此功能是 numpy 1.15.0 中的新功能。

https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.take_along_axis.html

请注意，data_arr 和 index_arr 需要具有相同的维数。因此，您需要将 index_array 重塑为 3 维，然后再将结果重塑为 2 维。IE：

start_time = time.time()
demos_arr = np.take_along_axis(img_arr, bayer_arr.reshape([img_height, img_width, 1]), axis=2).reshape([img_height, img_width])
end_time = time.time()

沿轴的计时结果与 ogrid 实现相同。

于 2018-12-07T11:20:40.370 回答