k-means 图像分割

经典的无监督聚类算法，不多说，上代码。

 import numpy as np
 import pandas as pd
 import copy
 import matplotlib.pyplot as plt
 
 pic = plt.imread('cs-nonoise.jpg')
 # plt.imshow(pic)
 # pic.shape   #(1200, 800)
 data = pic.reshape(-1, 3)
 
 def kmeans_wave(n, k, data):  # n为迭代次数， k为聚类数目， data为输入数据
     data_new = copy.deepcopy(data)
     data_new = np.column_stack((data_new, np.ones(1200*800)))   # 扩展一个维度用来存放标签
     center_point = np.random.choice(1200*800, k, replace=False)     # 随机选择初始点
     center = data_new[center_point,:]
     distance = [[] for i in range(k)]   # 距离度量
     for i in range(n):
          for j in range(k):
              distance[j] = np.sqrt(np.sum(np.square(data_new - np.array(center[j])), axis=1))     # 更新距离
          data_new[:,3] = np.argmin(np.array(distance), axis=0)   # 将最小距离的类别标签作为当前数据的类别
          for l in range(k):
              center[l] = np.mean(data_new[data_new[:,3]==l], axis=0)   # 更新聚类中心
 
     return data_new
 
 if __name__ == '__main__':
     data_new = kmeans_wave(100,5,data)
     print(data_new.shape)
     # data_new = np.delete(data_new, 3, axis=1)
     # print(data_new.shape)
     pic_new = data_new[:,3].reshape(1200,800)   # 将多个标签展示出来
     plt.imshow(pic_new)
     plt.show()

结果：

        

原图                                                                                                                            k=5 结果图