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 結果圖