如何有条件地叠加两个 numpy 数组/图像

Question

给定 2 个相同形状的图像/numpy 数组，例如：

• 数组A，形状=(100,100,3)。包含红色、蓝色和黑色像素。
• 数组B，形状=(100,100,3)。包含绿色、黄色和黑色像素。

现在我想将 arrayB 覆盖在 arrayA 上，条件是： （arrayC = 将 arrayB 覆盖在 arrayA 上）

1. 如果 arrayA 像素值为红色或蓝色，则 arrayC 像素 = arrayA 像素颜色。
2. 如果 arrayA 像素为黑色，则 arrayC 像素 = arrayB 像素颜色。

我怎样才能在 numpy 中做到这一点？

Answer 1

在 Numpy 中：

red = [1, 0, 0]
blue = [0, 0, 1]
black = [0, 0, 0]
shape = (5, 5, 3)
arrayA = np.zeros(shape)
arrayA[1, 1, :] = red
arrayA[2, 2, :] = blue
arrayB = np.ones(shape)
arrayC = np.full(shape, np.nan)

# Change selected pixels according to condition 1
pixels_to_change = np.where(
    (arrayA == red).all(axis=2) | (arrayA == blue).all(axis=2)
)
arrayC[pixels_to_change] = arrayA[pixels_to_change]

# Change selected pixels according to condition 2
pixels_to_change = np.where((arrayA == black).all(axis=2))
arrayC[pixels_to_change] = arrayB[pixels_to_change]

# Test pixels
assert np.array_equal(arrayC[1, 1], red)
assert np.array_equal(arrayC[2, 2], blue)
assert np.array_equal(arrayC[0, 0], arrayB[0, 0])