如何使用numpy进行多维数组查找?

时间:2016-03-05 20:38:31

标签: python numpy

我有这个RGB图像:

# 2D r, g and b channels
_r = np.array([[1,1],
               [0,0],
               [1,0]])

_g = np.array([[0,0],
               [1,1],
               [1,0]])

_b = np.array([[0,1],
               [1,0],
               [0,1]])

# Stack RGB into a single matrix
A = np.dstack((_r,_g,_b))

我想用这个图像来索引另一个数组:

X = np.array(
  [[[[ 0.,  0.,  0.],
     [ 0.,  0.,  0.]],

    [[.7,  0.,  0.],
     [ 0.,  0.,  0.]]],


   [[[ 0.,  0.,  0.],
     [ 0., .7,  0.]],

    [[ 0.,  0.,  .7],
     [ 0.,  0.,  0.]]]])

对于A中的每个rgb三元组,我想从X中抓取三个元素,然后我可以分别添加回r,g和b。

lookup = X[A]  # this is wrong
A += lookup

我尝试了很多东西,但无法获得正确的查找

A += X[A]             # nope
A += X[ zip(A) ]      # nope
A += X[ zip(*A.T) ]   # nope
A += X[ zip(A.T) ].T  # nope
...

添加后,A应为:

array([[[ 1. ,  0. ,  0. ],
        [ 1. ,  0.7,  1. ]],

       [[ 0. ,  1. ,  1. ],
        [ 0.7,  1. ,  0. ]],

       [[ 1. ,  1. ,  0.7],
        [ 0. ,  0. ,  1. ]]])

修改: 以下是我正在尝试进行的查找的更好描述:

# For each rgb triple at (x,y) in A...
r, g, b = A[x,y]
dr, dg, db = X[r,g,b]
A[x,y] += (dr, dg, db)

1 个答案:

答案 0 :(得分:1)

我认为你需要的只是

A += X[_r, _g, _b]

确保A实际上是一个浮点数组;否则你将无法获得预期的结果。 (您可以使用A.astype(float)将整数数组转换为浮点数组。)

如果您无法在实际代码中访问_r_g_b,则可以使用_r, _g, _b = A.transpose(2, 0, 1)来提取各个频道。