Python find colors in image array and fill 2D array with value

Question

I have an image with color segmentations and I want to find the pixels that have a color from a list of BGR colors. From those pixel indexes I want to fill a 2D array with some arbitrary value. I've accomplished this but it is terribly slow:

#load image with color segmentations
img = cv2.imread(segmented_img_path)
#create empty output array
output = np.zeros((img.shape[0], img.shape[1]))
#iterate over image and find colors
for i, row in enumerate(img):
    for j, bgr in enumerate(row):
        if np.any(np.all(np.isin(np.array(color_list),bgr,True),axis=1)):
            output[i,j] = float(some_value)

There's got to be a faster way of doing this, probably using np.where but I just can't figure it out.

Answer 1

I think this can be done as performed in the following example. Below is a simplified example that can be scaled to your needs.

m = np.array(([1,2,3], [4,5,6], [1,2,3]))
d = np.zeros((np.shape(m)))
BGR = [1,3]
for color in BGR:
   d[m==color] = color+1000

We simply loop through the color values you want to find in a BGR list and replace them in the for loop. Here color+1000 is an arbitrary value you refer to.

For your case it would appear as follows:

img = cv2.imread(segmented_img_path)
output = np.zeros((img.shape))
for bgr in BGR:
   output[img==bgr] = float(some_value)