numpy 元素条件格式化

Question

我有一个三维的numpy数组，想根据另一个元素的条件测试来改变某个特定的元素。（这个应用是为了调整RGBA图像数组中的'alpha'值，以便在3D的pyqtgraph图像中调整透明度——理想情况下，这个过程应该很快）。

a= np.ones((2,4,5),dtype=np.int) #create a ones array
a[0,0,0] = 3 #change a few values
a[0,2,0] = 3
a[1,2,0] = 3
print(a)
>>>[[[3 1 1 1 1]
  [1 1 1 1 1]
  [3 1 1 1 1]]

 [[1 1 1 1 1]
  [1 1 1 1 1]
  [3 1 1 1 1]]]

现在我想对最低维度的第一个元素进行条件测试，然后根据测试结果来改变最后一个元素。

if a[0,:,:] > 1:   #this does not work - for example only - if first element > 1
    a[3,:,:]=255   #then change the last element in the same dimension to 255

print(a) #this is my idealized result
>>>[[[3 1 1 1 255]
  [1 1 1 1 1]
  [3 1 1 1 255]]

 [[1 1 1 1 1]
  [1 1 1 1 1]
  [3 1 1 1 255]]]

Answer 1

这段代码看起来解决了问题：

mask = a[:,:,0] > 1

a[:,:,4][mask] = 255

所以索引的方式只需要稍微调整一下，然后就可以按照常规的方法应用一个掩码。

编辑 @Ophion 指出，这样写会更好：

a[mask,:,-1] = 255