使用scipy.cluster.hierarchy的linkage函数实现图像颜色聚类分割

嘿，我懂你想只靠颜色相似度做图像分割的需求，用层次聚类的思路完全靠谱！我来帮你把这个方案补全，顺便提些实用的细节，让代码能跑起来还能优化效果～

先给你补全并优化代码，再一步步拆解关键要点：

1. 完整可运行代码 #

假设你的img是已经加载好的(256,256,3) numpy数组，这里会处理像素值归一化、聚类执行和结果可视化：

import numpy as np
from scipy.cluster.hierarchy import linkage, fcluster
import matplotlib.pyplot as plt

# 处理像素值：把0-255的uint8转成0-1的float（避免数值范围干扰聚类）
img_normalized = img.astype(np.float32) / 255.0

# 将图像展平为[N, 3]的像素矩阵，N=256*256=65536
ppp = img_normalized.reshape(-1, 3)

# 执行层次聚类的linkage步骤——选ward方法最优，它会最小化类内方差，特别适合颜色聚类
# 注意：直接对65536个点聚类会很慢，后面会给优化方案
Z = linkage(ppp, method='ward', metric='euclidean')

# 指定聚类数量，比如分成5个颜色类
num_clusters = 5
labels = fcluster(Z, num_clusters, criterion='maxclust')

# 把聚类标签转回图像形状，得到分割结果
segmented_img = labels.reshape(img.shape[0], img.shape[1])

# 对比可视化原图和分割图
plt.figure(figsize=(10,5))
plt.subplot(121)
plt.imshow(img)
plt.title('Original Image')
plt.axis('off')

plt.subplot(122)
plt.imshow(segmented_img, cmap='viridis')
plt.title(f'Segmented Image ({num_clusters} color groups)')
plt.axis('off')
plt.show()

2. 关键优化与细节说明 #

• 性能救星：采样聚类
  直接对65536个像素做层次聚类（时间复杂度O(n²)）会非常慢，你可以先随机采样部分像素聚类，再用KNN给剩余像素贴标签，速度提升几十倍，效果几乎没差：

  # 随机采样1000个像素做聚类
  sample_indices = np.random.choice(ppp.shape[0], size=1000, replace=False)
  sample_points = ppp[sample_indices]
  Z_sample = linkage(sample_points, method='ward', metric='euclidean')
  sample_labels = fcluster(Z_sample, num_clusters, criterion='maxclust')
  
  # 用KNN给所有像素分配标签
  from sklearn.neighbors import KNeighborsClassifier
  knn = KNeighborsClassifier(n_neighbors=5)
  knn.fit(sample_points, sample_labels)
  labels = knn.predict(ppp)
  

• 聚类数量怎么选？
  不想手动指定num_clusters的话，可以画树状图找聚类分界点——看树状图里距离突然跳变的“断层”，就是合适的聚类数量：

  plt.figure(figsize=(15,5))
  dendrogram(Z_sample)
  plt.title('Dendrogram of Color Clusters')
  plt.xlabel('Sample Pixels')
  plt.ylabel('Euclidean Distance')
  plt.show()
  

• 更贴合人眼的颜色聚类
  RGB空间的欧氏距离和人眼感知的颜色相似度不完全匹配，你可以把图像转成Lab颜色空间再聚类，结果会更符合视觉直觉：

  from skimage.color import rgb2lab
  img_lab = rgb2lab(img_normalized)
  ppp_lab = img_lab.reshape(-1, 3)
  # 后续聚类步骤用ppp_lab代替ppp即可
  

3. 必踩坑提醒 #

• 一定要把uint8类型的0-255像素值转成0-1的float，不然大数值的RGB分量会主导聚类结果，完全打乱颜色相似性的判断。

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