如何修改K-means聚类图像配色，使其匹配指定色彩？

我期望得到的效果：

但实际生成的图像是这样的：

我希望最终图像能包含红、蓝、绿、青、黄这些颜色，参考效果：

以下是我的代码：

import numpy as np
import cv2 as cv

def kmeans_clustering(image, K):
    pixels = image.reshape((-1, 3)).astype(np.float32)
    pixels_normalized = pixels / 255.0
    np.random.seed(0)
    centers = np.random.rand(K, 3)
    iterations = 10
    for _ in range(iterations):
        
        distances = np.linalg.norm(pixels_normalized[:, np.newaxis] - centers, axis=2)
        labels = np.argmin(distances, axis=1)
        for i in range(K):
            cluster_pixels = pixels_normalized[labels == i]
            if len(cluster_pixels) > 0:
                centers[i] = np.mean(cluster_pixels, axis=0)
            
    centers_rgb = np.fliplr((centers * 255).astype(np.uint8))    
    segmented_image = centers_rgb[labels].reshape(image.shape)

    return segmented_image

image_rgb = cv.imread('sunflower.png')
image_rgb = cv.cvtColor(image_rgb, cv.COLOR_BGR2RGB)

k_values = [2, 4, 6, 8]
segmented_images = []

for k in k_values:
    segmented_image = kmeans_clustering(image_rgb, k)
    segmented_images.append(segmented_image)

stacked_image = np.vstack((np.hstack((segmented_images[0], segmented_images[1])),
                           np.hstack((segmented_images[2], segmented_images[3]))))

cv.imshow('Segmented Image', stacked_image)
cv.waitKey(0)
cv.destroyAllWindows()

你的代码目前是用聚类计算出的平均颜色作为分割色块，所以会出现偏灰暗的色调。要改成指定的鲜明色彩，核心是把聚类结果映射到你预设的目标颜色上，步骤如下：

1. 定义目标色彩集合
   先把需要的颜色按RGB格式列出，K值大于5时补充额外颜色保证数量足够：

target_colors = np.array([
    [255, 0, 0],    # 红
    [0, 0, 255],    # 蓝
    [0, 255, 0],    # 绿
    [0, 255, 255],  # 青
    [255, 255, 0],  # 黄
    [255, 165, 0],  # 橙（适配K=6）
    [128, 0, 128],  # 紫（适配K=8）
    [255, 255, 255] # 白（适配K=8）
], dtype=np.uint8)

2. 修改聚类函数的颜色映射逻辑
   保留原有的聚类分组逻辑，但把聚类得到的平均中心替换成最接近的目标颜色：

def kmeans_clustering(image, K, target_colors):
    pixels = image.reshape((-1, 3)).astype(np.float32)
    pixels_normalized = pixels / 255.0
    np.random.seed(0)
    centers = np.random.rand(K, 3)
    iterations = 10
    for _ in range(iterations):
        
        distances = np.linalg.norm(pixels_normalized[:, np.newaxis] - centers, axis=2)
        labels = np.argmin(distances, axis=1)
        for i in range(K):
            cluster_pixels = pixels_normalized[labels == i]
            if len(cluster_pixels) > 0:
                centers[i] = np.mean(cluster_pixels, axis=0)
    
    # 将聚类中心转为0-255的RGB格式
    centers_rgb = (centers * 255).astype(np.uint8)
    
    # 为每个聚类匹配最接近的目标颜色
    mapped_colors = []
    for center in centers_rgb:
        color_distances = np.linalg.norm(target_colors - center, axis=1)
        closest_color = target_colors[np.argmin(color_distances)]
        mapped_colors.append(closest_color)
    mapped_colors = np.array(mapped_colors)
    
    # 生成最终分割图像
    segmented_image = mapped_colors[labels].reshape(image.shape)

    return segmented_image

3. 调整主代码的函数调用
   把目标颜色集合传入聚类函数：

image_rgb = cv.imread('sunflower.png')
image_rgb = cv.cvtColor(image_rgb, cv.COLOR_BGR2RGB)

k_values = [2, 4, 6, 8]
segmented_images = []

# 定义目标颜色集合
target_colors = np.array([
    [255, 0, 0],    # 红
    [0, 0, 255],    # 蓝
    [0, 255, 0],    # 绿
    [0, 255, 255],  # 青
    [255, 255, 0],  # 黄
    [255, 165, 0],  # 橙
    [128, 0, 128],  # 紫
    [255, 255, 255] # 白
], dtype=np.uint8)

for k in k_values:
    segmented_image = kmeans_clustering(image_rgb, k, target_colors)
    segmented_images.append(segmented_image)

stacked_image = np.vstack((np.hstack((segmented_images[0], segmented_images[1])),
                           np.hstack((segmented_images[2], segmented_images[3]))))

cv.imshow('Segmented Image', stacked_image)
cv.waitKey(0)
cv.destroyAllWindows()

4. 额外说明

• 如果希望K>5时也只使用指定的5种颜色，可以删除补充的颜色，让多个聚类中心自动映射到最接近的5种颜色中，不过会降低分割的区分度。
• 若要更精准的颜色匹配，可以调整距离计算的方式（比如改用HSV色彩空间计算距离，比RGB更符合人眼感知）。

内容的提问来源于stack exchange，提问作者ERICA