Разделить изображение на N изображений, где N — количество цветов, появляющихся на нем.

Я пытаюсь разделить изображение в зависимости от содержащихся в нем цветов.

Мои предыдущие шаги заключались в том, чтобы упростить его до 3 цветов, используя алгоритм KMeans, который предлагает Sklearn, и я получил результат, подобный следующему изображению.

Теперь мне нужно разделить его на 3 изображения по одному для каждого цвета. И получить что-то похожее на это (я сделал это с фотошопом).

Примеры выполнены в черно-белом цвете, потому что на самом деле, если я могу разделить, мне больше не нужны цвета. Но я бы сделал то же самое с 3-х цветными изображениями.

Маска 1:

Маска 2:

Маска 3:

Я нашел этот вопрос, но не могу достичь своей цели.

Я думал о разделении по каналам, но думаю, что это неправильно.

# 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

Я думаю, что ключ в моем алгоритме kmeans, потому что с его помощью я получаю labels и centroids своих цветов, но я действительно не знаю, как это сделать, и я не могу найти никого, кто бы это делал.

Мой алгоритм KMeans:

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

Может ли кто-нибудь помочь мне, пожалуйста?

Большое спасибо!

Обновлено: Из ответа Хихикомори.

Из ответа Хихикомори я понимаю, что я должен сделать следующее, это основано на вопросе, который я связал ранее, но проблема в том, что я получаю 3 черные маски без какого-либо контура, поэтому я подумал, что это мне не подойдет.

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