k-mean算法进行图像分割

import numpy as np
# import matplotlib.pyplot as plt

import cv2
def knn(data, iter, k):

    data = cv2.merge([data, data, data])
    data = data.reshape(-1, 3)
    data = np.column_stack((data, np.ones(row*col)))
   
    cluster_center = data[np.random.choice(row*col, k)]
    
    distance = [[] for i in range(k)]
    for i in range(iter):
    
        for j in range(k):
            distance[j] = np.sqrt(np.sum((data - cluster_center[j])**2, axis=1))
        
        data[:, 3] = np.argmin(distance, axis=0)
        
        for j in range(k):
            cluster_center[j] = np.mean(data[data[:, 3] == j], axis=0)
    return data[:, 3]
 
 
if __name__ == "__main__":
    img = cv2.imread('5.jpg', 0)
    row = img.shape[0]
    col = img.shape[1]
    # cv2.subplot(121)
    # cv2.imshow("img", img)
    # cv2.waitKey(0)
    image_show = knn(img, 10, 2)
    image_show = image_show.reshape(row, col)
    # cv2.subplot(122)
    cv2.imshow("result", image_show)
    cv2.waitKey(0)

参考：基于K-means的图像分割（python代码详解）