如何将两个numpy数组加为一个？

Question

我有一个数组preds，其值为0,1 or 2。

preds的形状和值看起来像这样：

torch.Size([1, 256, 256])
tensor([[[1, 1, 0,  ..., 1, 1, 2],
         [1, 0, 1,  ..., 1, 1, 1],
         [0, 0, 0,  ..., 1, 1, 1],
         ...,
         [0, 1, 1,  ..., 1, 1, 1],
         [1, 1, 1,  ..., 1, 1, 1],
         [1, 1, 1,  ..., 1, 1, 1]]])

我只关心第二维和第三维（256 x 256）。我想使用matplotlib创建一个遮罩，该遮罩可分别绘制0，1和2。基本上，我想可视化图像中具有0的所有位置，具有1的位置以及具有2的位置。我要为此划算。

我在想类似的东西。例如对于类1：

1. 获取我们有1个地方的地点并为其赋值1，并为其他索引分配值0。
2. 使用plt.show绘制

示例：

zero = (preds==0).nonzero() #index of all values which have 0 
nonzero = (preds!=0).nonzero()    # index of all value which are not 0
# assign 1 to all zero
# assign 0 to all nonzero
# join the above two array
# plot the array

我只是不知道如何在我的伪代码中添加这两个数组（nonzero和zero）。

是否有更Python化的方法？我不想遍历所有索引值，然后为它们分配值。

Answer 1

我能够按照评论中的建议使用np.where完成此操作。看起来像这样：

#preds in the matrix where I want the 0 and 1s
preds_clear = np.where(preds==0, 1, 0)
plt.imshow(preds_clear[0,:,:])