使用基于另一个维度的np.where设置numpy切片

Posted 2021-04-27

如何根据第4个通道中的值设置4通道numpy阵列的前3个通道中的值？是否可以将numpy切片作为l值来实现？

给出一个3乘2像素的numpy数组，有4个通道

a = np.arange(24).reshape(3,2,4)
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],
       [[ 8,  9, 10, 11],
        [12, 13, 14, 15]],
       [[16, 17, 18, 19],
        [20, 21, 22, 23]]])

我可以选择第4个通道为模3的切片。

px = np.where(0==a[:,:,3]%3)
(array([0, 1], dtype=int64), array([0, 1], dtype=int64))

a[px]
array([[ 0,  1,  2,  3],
       [12, 13, 14, 15]])

现在我想将a中这些行中的前3个通道设置为0，这样结果如下所示：

a
array([[[ 0,  0,  0,  3],
        [ 4,  5,  6,  7]],
       [[ 8,  9, 10, 11],
        [ 0,  0,  0, 15]],
       [[16, 17, 18, 19],
        [20, 21, 22, 23]]])

我试过了

a[px][:,0:3] = 0

但这使得数组保持不变。

我读了Setting values in a numpy arrays indexed by a slice and two boolean arrays并且不明白如何使用布尔索引来设置前3个通道。