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

Question

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

Answer 1

You can find the unique colours in your image with np.unique() and then iterate over them setting each pixel to either white or black depending whether it is equal to that colour or not:

#!/usr/bin/env python3

import cv2
import numpy as np

# Load image
im = cv2.imread('cheese.png')

# Reshape into a tall column of pixels, each with 3 RGB pixels and get unique rows (colours)
colours  = np.unique(im.reshape(-1,3), axis=0)

# Iterate over the colours we found
for i,colour in enumerate(colours):
    print(f'DEBUG: colour {i}: {colour}')
    res = np.where((im==colour).all(axis=-1),255,0)
    cv2.imwrite(f'colour-{i}.png', res)

Sample Output

DEBUG: colour 0: [0 0 0]
DEBUG: colour 1: [0 141 196]
DEBUG: colour 2: [1 102 133]

Answer 2

color1 = (0,0,160)
color2 = (0,160,160)
color3 = (160,160,160)
img = np.zeros((640,480,3),np.uint8)

img[100:200,100:200] = color1
img[150:250,150:250] = color2
img[200:300,200:300] = color3

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

cv2.imshow('origin', img)
cv2.imshow('img1', img1)
cv2.imshow('img2', img2)
cv2.imshow('img3', img3)
cv2.waitKey(0)