Replace vectors in 2d numpy array

Question

I have a numpy array of shape (height, width, 3) loaded from an image. I want to replace all black pixels [0, 0, 0] with a specific color [r, g, b]. Is there a numpy way to do this?

Answer 1

import numpy as np
orig_color = (0, 0, 0)
replacement_color = (r, g, b)

data[(data == orig_color).all(axis = -1)] = replacement_color