Select 并在 numpy n 维内部数组上应用条件

Select and apply condition on numpy n-dimensional inner array

假设我有一个形状为 (2, 4, 3) 的 numpy 数组,如下所示:

import numpy as np

arr = np.array([[[248,  26,   4],
                 [ 99, 126, 156],
                 [ 80, 240, 232],
                 [136,  27, 216]],

                [[221, 130, 119],
                 [253, 188, 232],
                 [159,  21,  98],
                 [ 12,  35,  50]]])

我想检查较小数组中的每一项是否满足条件,例如:

conditions_list = [(arr[:,:,0] > 100 and arr[:,:,1] > 100 and arr[:,:,2] > 100),
                   (arr[:,:,0] < 25 and arr[:,:,1] < 50 and arr[:,:,2] < 50)]
choices_list = [(arr[:,:,0] = 255, arr[:,:,1] = 255, arr[:,:,2] = 255),
                (arr[:,:,0] = 0, arr[:,:,1] = 0, arr[:,:,2] = 0)]
new_array = np.select(conditions_list, choices, default=(arr[:,:,0], arr[:,:,1], arr[:,:,2]))
#This code doesn't work, but it represents the idea of what I need to get.

预期结果如下:

([[[248,  26,   4],
   [ 99, 126, 156],
   [ 80, 240, 232],
   [136,  27, 216]],

  [[255, 255, 255],
   [255, 255, 255],
   [159,  21,  98],
   [  0,   0,   0]]])

当我 运行 上面的代码时,我得到一个异常,说我应该使用 arr.any()arr.all() 而不是 conditions_list,但是这些功能不符合我需要的。

如何尽可能高效地根据条件修改内部数组?可以有数以千计的数组最多可以整形为 (3000, 3000, 3),这就是为什么我认为 np.select() 是一个不错的选择。


编辑

我知道我可以使用定义自定义函数的列表理解,但这太慢了,因为它会遍历每个项目。可以这样做:

new_arr = [[cond_func(x) for x in y] for y in arr]

def cond_func(x):
    if x[0] > 100 and x[1] > 100 and x[2] > 100:
        x[0], x[1], x[2] = 255, 255, 255
    elif x[0] < 25 and x[1] < 50 and x[2] < 50:
        x[0], x[1], x[2] = 0, 0 ,0
    return(x)
In [78]: arr = np.array([[[248,  26,   4],
    ...:                  [ 99, 126, 156],
    ...:                  [ 80, 240, 232],
    ...:                  [136,  27, 216]],
    ...: 
    ...:                 [[221, 130, 119],
    ...:                  [253, 188, 232],
    ...:                  [159,  21,  98],
    ...:                  [ 12,  35,  50]]])
    ...: 

这是处理 >100 案例的 'masking' 方法:

In [79]: (arr>100)
Out[79]: 
array([[[ True, False, False],
        [False,  True,  True],
        [False,  True,  True],
        [ True, False,  True]],

       [[ True,  True,  True],
        [ True,  True,  True],
        [ True, False, False],
        [False, False, False]]])
In [80]: (arr>100).all(axis=2)
Out[80]: 
array([[False, False, False, False],
       [ True,  True, False, False]])
In [81]: arr[_]
Out[81]: 
array([[221, 130, 119],
       [253, 188, 232]])
In [82]: arr1=arr.copy()
In [83]: arr1[(arr>100).all(axis=2)]=255
In [84]: arr1
Out[84]: 
array([[[248,  26,   4],
        [ 99, 126, 156],
        [ 80, 240, 232],
        [136,  27, 216]],

       [[255, 255, 255],
        [255, 255, 255],
        [159,  21,  98],
        [ 12,  35,  50]]])

对于第二个条件,arr[:,:,2]:

使用 75
In [94]: arr1[(arr<np.array([25,50,75])).all(axis=2)]=0
In [95]: arr1
Out[95]: 
array([[[248,  26,   4],
        [ 99, 126, 156],
        [ 80, 240, 232],
        [136,  27, 216]],

       [[255, 255, 255],
        [255, 255, 255],
        [159,  21,  98],
        [  0,   0,   0]]])

下面是将这些掩码与 select 一起使用的方法。请注意,我必须调整蒙版的尺寸,以便它们与 3d default.

一起使用
In [104]: mask1 = (arr>100).all(axis=2)
In [105]: mask2 = (arr<np.array([25,50,75])).all(axis=2)
In [106]: mask1.shape
Out[106]: (2, 4)
In [107]: arr.shape
Out[107]: (2, 4, 3)
In [108]: np.select([mask1[:,:,None],mask2[:,:,None]],[255,0], default=arr)
Out[108]: 
array([[[248,  26,   4],
        [ 99, 126, 156],
        [ 80, 240, 232],
        [136,  27, 216]],

       [[255, 255, 255],
        [255, 255, 255],
        [159,  21,  98],
        [  0,   0,   0]]])

或使用:

In [109]: mask1 = (arr>100).all(axis=2, keepdims=True)
In [110]: mask1.shape
Out[110]: (2, 4, 1)

要使用您的语法,我们需要添加 () 并使用 & 而不是 and:

In [111]: (arr[:,:,0] > 100) & (arr[:,:,1] > 100) & (arr[:,:,2] > 100)
Out[111]: 
array([[False, False, False, False],
       [ True,  True, False, False]])