Split image in N images, where N is the number of colors appearing on it – Photoshop

I’m trying to split an image depending on the colors it contains.

My previous steps have been to simplify it to just 3 colors using the KMeans algorithm that Sklearn offers and I get a result like the following image.

Now I need to split it in 3 images one for each color. And obtain something similar to this (I have done it with photoshop).

The examples are done in black and white because really if I can do the division, I no longer need the colors. But I would do just as well with 3 color images.

Mask 1:

Mask 2:

Mask 3:

I found this question, but I can’t achieve my goal.

I have thought about separating by channels, but I think it is wrong.

# set green and red channels to 0
blue_img[:, :, 1] = 0
blue_img[:, :, 2] = 0
# set blue and red channels to 0
green_img[:, :, 0] = 0
green_img[:, :, 2] = 0
# set blue and green channels to 0
red_img[:, :, 0] = 0
red_img[:, :, 1] = 0

I think the key is in my kmeans algorithm, because with it I obtain labels and centroids of my colors but I really don’t know how to do it, and I can’t find anybody doing it.

My KMeans algorithm is:

def get_colors(img, number_of_colors, show_chart, show_segmented_img):
    
    modified_image = img.reshape(img.shape[0]*img.shape[1], 3)
    
    myKMeans = KMeans(n_clusters = number_of_colors)
    
    labels = myKMeans.fit_predict(modified_image)
    
    counts = Counter(labels)
    
    centroids = myKMeans.cluster_centers_
    
    ordered_colors = [centroids[i] for i in counts.keys()]
    
    hex_colors = [RGB2HEX(ordered_colors[i]) for i in counts.keys()]
    
    rgb_colors = [ordered_colors[i] for i in counts.keys()]
    
    if (show_chart):
        plt.figure(figsize = (8, 6))
        plt.pie(counts.values(), labels = hex_colors, colors = hex_colors)
        plt.show()

    if (show_segmented_img):
        centroids = np.uint8(centroids)
        segmented_data = centroids[labels.flatten()]
        segmented_image = segmented_data.reshape(img.shape)
        segmented_image = cv2.cvtColor(segmented_image, cv2.COLOR_RGB2BGR)
        cv2.imwrite('segmentedImg.png', segmented_image)

    return hex_colors, rgb_colors

can somebdoy help me please?

Thank you very much!

EDIT: From Hihikomori’s answer.

From Hihikomori’s answer, I understand that I should do the following, this is based on the question that I linked before, but the problem is that I get 3 black masks without any contour so I thought that this would not suit me.

def get_colors(img, number_of_colors, show_chart, show_segmented_img):
    
    modified_image = img.reshape(img.shape[0]*img.shape[1], 3)
    
    myKMeans = KMeans(n_clusters = number_of_colors)
    
    labels = myKMeans.fit_predict(modified_image)
    
    counts = Counter(labels)
    
    centroids = myKMeans.cluster_centers_
    
    ordered_colors = [centroids[i] for i in counts.keys()]
    
    hex_colors = [RGB2HEX(ordered_colors[i]) for i in counts.keys()]
    
    rgb_colors = [ordered_colors[i] for i in counts.keys()]

    # TRYING THE ASNWER
    color1, color2,color3 = rgb_colors
    first_color_indices = np.where(np.all(img == color1, axis=-1))
    second_color_indices = np.where(np.all(img == color2, axis=-1))
    third_color_indices = np.where(np.all(img == color3, axis=-1))

    img1 = np.zeros_like(img)
    img1[first_color_indices]=color1

    img2 = np.zeros_like(img)
    img2[second_color_indices]=color2

    img3 = np.zeros_like(img)
    img3[third_color_indices]=color3

    print('***')
    cv2_imshow(img1)
    print('***')
    cv2_imshow(img2)
    print('***')
    cv2_imshow(img3)
    print('***')
    
    if (show_chart):
        plt.figure(figsize = (8, 6))
        plt.pie(counts.values(), labels = hex_colors, colors = hex_colors)
        plt.show()

    if (show_segmented_img):
        centroids = np.uint8(centroids)
        segmented_data = centroids[labels.flatten()]
        segmented_image = segmented_data.reshape(img.shape)
        segmented_image = cv2.cvtColor(segmented_image, cv2.COLOR_RGB2BGR)
        cv2.imwrite('segmentedImg.png', segmented_image)

    return hex_colors, rgb_colors