Numpy Image Array 用 1 替换特定颜色，用 0 替换其他颜色

Question

我必须对图像进行语义分割编码。输入图像的形状为 (128, 256, 3)，具有 128 x 256 RGB 值。我想要一个 (128, 256) 的输出形状，其中每 1 表示像素与给定颜色匹配，0 表示存在另一个 RGB 值。

[[20, 20, 20], [30, 30, 30], [40, 40, 40]] 和过滤器 [20,20,20] 应该导致 [1, 0, 0]

任何帮助将不胜感激。

此方法最好适用于包含 16 张图片的形状数组 (16, 128, 256, 3)，并将过滤器应用于每张图片。

Answer 1

对一张图片使用np.where：

filter_pixel = np.array([20, 20, 20])
image = np.array([[[20, 20, 20], [30, 30, 30], [40, 40, 40]],
                  [[20, 10, 20], [20, 20, 20], [40, 40, 40]]])

new_image = np.where(np.all(image == filter_pixel, axis=2), 1,0)
print(new_image)

输出：

[[1 0 0]
 [0 1 0]]

编辑，多张图片：

filter_pixels = np.array([[20, 20, 20], [30, 30, 30]])
filter_pixels = filter_pixels[:, np.newaxis, np.newaxis, :]
images = np.array([[[[20, 20, 20], [30, 30, 30], [40, 40, 40]],
                    [[20, 10, 20], [20, 20, 20], [40, 40, 40]]],
                  [[[20, 20, 20], [30, 30, 30], [40, 40, 40]],
                   [[20, 10, 30], [20, 20, 20], [30, 30, 30]]]])

new_images = np.where(np.all(images == filter_pixels, axis=3), 1, 0)

print(new_images)

输出：

[[[1 0 0]
  [0 1 0]]

 [[0 1 0]
  [0 0 1]]]