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

Question

如何根据第 4 个 channel 中的值设置 4 channel numpy 数组的前 3 个 channel 中的值？是否可以使用 numpy 切片作为左值来做到这一点？

给定一个具有 4 个 channel 的 3 x 2 像素 numpy 数组

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 个 channel 为模 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]])

现在我想将这些行中的前 3 个 channel 设置为 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并且不明白如何使用 bool 索引仅设置前 3 个 channel 。

Answer 1

这是一种方法:

>>> px0, px1 = np.where(0==a[:,:,3]%3)
>>> a[px0, px1, :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]]])

或者

>>> px = np.where(0==a[:,:,3]%3)
>>> a[..., :3][px] = 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, np.s_[: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]]])